"""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.utilities.blas import axpy
from gpaw.utilities import unpack
from gpaw.eigensolvers.eigensolver import Eigensolver


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(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(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(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(float(np.real(b)))
                c = comm.sum(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(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
