"""Module defining  ``Eigensolver`` classes."""

from math import pi, sqrt, sin, cos, atan2

import numpy as np
from numpy import dot
from ase.units import Hartree

from gpaw.eigensolvers.eigensolver import Eigensolver
from gpaw.utilities import unpack_hermitian
from gpaw.utilities.blas import axpy


class CG(Eigensolver):
    """Conjugate gardient eigensolver

    It is expected that the trial wave functions are orthonormal
    and the integrals of projector functions and wave functions
    are already calculated.

    Solution steps are:

    * Subspace diagonalization
    * Calculate all residuals
    * Conjugate gradient steps
    """

    def __init__(self, niter=4, rtol=0.30, tw_coeff=False):
        """Construct conjugate gradient eigen solver.

        parameters:

        niter: int
            Maximum number of conjugate gradient iterations per band
        rtol: float
            If change in residual is less than rtol, iteration for band is
            not continued

        """
        Eigensolver.__init__(self)
        self.niter = niter
        self.rtol = rtol
        self.orthonormalization_required = False
        self.tw_coeff = tw_coeff

        self.tolerance = None

    def __repr__(self):
        return 'CG(niter=%d, rtol=%5.1e)' % (self.niter, self.rtol)

    def todict(self):
        return {'name': 'cg', 'niter': self.niter}

    def initialize(self, wfs):
        if wfs.bd.comm.size > 1:
            raise ValueError('CG eigensolver does not support band '
                             'parallelization.  This calculation parallelizes '
                             'over %d band groups.' % wfs.bd.comm.size)
        if wfs.mode == 'pw' and wfs.gd.comm.size > 1:
            raise ValueError('CG eigensolver does not support domain '
                             'parallelization in PW-mode.')
        Eigensolver.initialize(self, wfs)

    def iterate_one_k_point(self, ham, wfs, kpt, weights):
        """Do conjugate gradient iterations for the k-point"""
        self.timer.start('CG')

        niter = self.niter

        phi_G, phi_old_G, Htphi_G = wfs.empty(3, q=kpt.q)

        comm = wfs.gd.comm
        if self.tw_coeff:
            # Wait!  What business does the eigensolver have changing
            # the properties of the Hamiltonian?  We are not updating
            # the Hamiltonian here.  Moreover, what is supposed to
            # happen if this function is called multiple times per
            # iteration?  Then we keep dividing the potential by the
            # same number.  What on earth is the meaning of this?
            #
            # Also the parameter tw_coeff is undocumented.  What is it?
            ham.vt_sG /= self.tw_coeff
            # Assuming the ordering in dH_asp and wfs is the same
            for a in ham.dH_asp.keys():
                ham.dH_asp[a] /= self.tw_coeff

        psit = kpt.psit
        R = psit.new(buf=wfs.work_array)
        P = kpt.projections
        P2 = P.new()

        self.subspace_diagonalize(ham, wfs, kpt)

        Htpsit = psit.new(buf=self.Htpsit_nG)

        R.array[:] = Htpsit.array
        self.calculate_residuals(kpt, wfs, ham, psit,
                                 P, kpt.eps_n, R, P2)

        total_error = 0.0
        for n in range(self.nbands):
            N = self.nbands
            R_G = R.array[n]
            Htpsit_G = Htpsit.array[n]
            psit_G = psit.array[n]
            gamma_old = 1.0
            phi_old_G[:] = 0.0
            error = np.real(wfs.integrate(R_G, R_G))
            for nit in range(niter):
                if (error * Hartree**2 < self.tolerance / self.nbands):
                    break

                ekin = self.preconditioner.calculate_kinetic_energy(psit_G,
                                                                    kpt)

                pR_G = self.preconditioner(R_G, kpt, ekin)

                # New search direction
                gamma = comm.sum_scalar(np.vdot(pR_G, R_G).real)
                phi_G[:] = -pR_G - gamma / gamma_old * phi_old_G
                gamma_old = gamma
                phi_old_G[:] = phi_G[:]

                # Calculate projections
                P2_ai = wfs.pt.dict()
                wfs.pt.integrate(phi_G, P2_ai, kpt.q)

                # Orthonormalize phi_G to all bands
                self.timer.start('CG: orthonormalize')
                self.timer.start('CG: overlap')
                overlap_n = wfs.integrate(psit.array[:N], phi_G,
                                          global_integral=False)
                self.timer.stop('CG: overlap')
                self.timer.start('CG: overlap2')
                for a, P2_i in P2_ai.items():
                    P_ni = kpt.P_ani[a]
                    dO_ii = wfs.setups[a].dO_ii
                    overlap_n += np.dot(P_ni[:N].conjugate(),
                                        np.dot(dO_ii, P2_i))
                self.timer.stop('CG: overlap2')
                comm.sum(overlap_n)

                phi_G -= psit.array[:N].T.dot(overlap_n).T

                for a, P2_i in P2_ai.items():
                    P_ni = kpt.P_ani[a]
                    P2_i -= np.dot(overlap_n, P_ni[:N])

                norm = wfs.integrate(phi_G, phi_G, global_integral=False)
                for a, P2_i in P2_ai.items():
                    dO_ii = wfs.setups[a].dO_ii
                    norm += np.vdot(P2_i, np.dot(dO_ii, P2_i))
                norm = comm.sum_scalar(float(np.real(norm)))
                phi_G /= sqrt(norm)
                for P2_i in P2_ai.values():
                    P2_i /= sqrt(norm)
                self.timer.stop('CG: orthonormalize')

                # find optimum linear combination of psit_G and phi_G
                an = kpt.eps_n[n]
                wfs.apply_pseudo_hamiltonian(kpt, ham,
                                             phi_G.reshape((1,) + phi_G.shape),
                                             Htphi_G.reshape((1,) +
                                                             Htphi_G.shape))
                b = wfs.integrate(phi_G, Htpsit_G, global_integral=False)
                c = wfs.integrate(phi_G, Htphi_G, global_integral=False)
                for a, P2_i in P2_ai.items():
                    P_i = kpt.P_ani[a][n]
                    dH_ii = unpack_hermitian(ham.dH_asp[a][kpt.s])
                    b += dot(P2_i, dot(dH_ii, P_i.conj()))
                    c += dot(P2_i, dot(dH_ii, P2_i.conj()))
                b = comm.sum_scalar(float(np.real(b)))
                c = comm.sum_scalar(float(np.real(c)))

                theta = 0.5 * atan2(2 * b, an - c)
                enew = (an * cos(theta)**2 +
                        c * sin(theta)**2 +
                        b * sin(2.0 * theta))
                # theta can correspond either minimum or maximum
                if (enew - kpt.eps_n[n]) > 0.0:  # we were at maximum
                    theta += pi / 2.0
                    enew = (an * cos(theta)**2 +
                            c * sin(theta)**2 +
                            b * sin(2.0 * theta))

                kpt.eps_n[n] = enew
                psit_G *= cos(theta)
                # kpt.psit_nG[n] += sin(theta) * phi_G
                axpy(sin(theta), phi_G, psit_G)
                for a, P2_i in P2_ai.items():
                    P_i = kpt.P_ani[a][n]
                    P_i *= cos(theta)
                    P_i += sin(theta) * P2_i

                if nit < niter - 1:
                    Htpsit_G *= cos(theta)
                    # Htpsit_G += sin(theta) * Htphi_G
                    axpy(sin(theta), Htphi_G, Htpsit_G)
                    # adjust residuals
                    R_G[:] = Htpsit_G - kpt.eps_n[n] * psit_G

                    coef_ai = wfs.pt.dict()
                    for a, coef_i in coef_ai.items():
                        P_i = kpt.P_ani[a][n]
                        dO_ii = wfs.setups[a].dO_ii
                        dH_ii = unpack_hermitian(ham.dH_asp[a][kpt.s])
                        coef_i[:] = (dot(P_i, dH_ii) -
                                     dot(P_i * kpt.eps_n[n], dO_ii))
                    wfs.pt.add(R_G, coef_ai, kpt.q)
                    error_new = np.real(wfs.integrate(R_G, R_G))
                    if error_new / error < self.rtol:
                        # print >> self.f, "cg:iters", n, nit+1
                        break
                    if (self.nbands_converge == 'occupied' and
                        kpt.f_n is not None and kpt.f_n[n] == 0.0):
                        # print >> self.f, "cg:iters", n, nit+1
                        break
                    error = error_new

            total_error += weights[n] * error
            # if nit == 3:
            #   print >> self.f, "cg:iters", n, nit+1
        if self.tw_coeff:  # undo the scaling for calculating energies
            for i in range(len(kpt.eps_n)):
                kpt.eps_n[i] *= self.tw_coeff
            ham.vt_sG *= self.tw_coeff
            # Assuming the ordering in dH_asp and wfs is the same
            for a in ham.dH_asp.keys():
                ham.dH_asp[a] *= self.tw_coeff

        self.timer.stop('CG')
        return total_error
