import numpy as np
import numpy.linalg as la
from ase.units import Hartree

from gpaw.lcao.overlap import NewTwoCenterIntegrals as TwoCenterIntegrals
from gpaw.lfc import BasisFunctions
from gpaw.utilities import unpack_hermitian
from gpaw.utilities.tools import dagger, lowdin, tri2full
from gpaw.lcao.tools import basis_subset2


def dots(*args):
    x = args[0]
    for M in args[1:]:
        x = np.dot(x, M)
    return x


def normalize(U, U2=None, norms=None):
    if norms is None:
        norms2 = np.ones(U.shape[1])
    else:
        norms2 = np.sqrt(norms)

    if U2 is None:
        for n, col in zip(norms2, U.T):
            col *= n / la.norm(col)
    else:
        for n, col1, col2 in zip(norms2, U.T, U2.T):
            norm = np.sqrt(np.vdot(col1, col1) + np.vdot(col2, col2))
            col1 *= n / norm
            col2 *= n / norm


def normalize2(C, S):
    C /= np.sqrt(dots(dagger(C), S, C).diagonal())


def get_rot(F_MM, V_oM, L):
    eps_M, U_MM = la.eigh(F_MM)
    indices = eps_M.real.argsort()[-L:]
    U_Ml = U_MM[:, indices]
    U_Ml /= np.sqrt(dots(U_Ml.T.conj(), F_MM, U_Ml).diagonal())

    U_ow = V_oM.copy()
    U_lw = np.dot(U_Ml.T.conj(), F_MM)
    for col1, col2 in zip(U_ow.T, U_lw.T):
        norm = np.sqrt(np.vdot(col1, col1) + np.vdot(col2, col2))
        col1 /= norm
        col2 /= norm
    return U_ow, U_lw, U_Ml


def condition_number(S):
    eps = la.eigvalsh(S).real
    return eps.max() / eps.min()


def eigvals(H, S):
    return np.sort(la.eigvals(la.solve(S, H)).real)


def get_bfs(calc):
    wfs = calc.wfs
    bfs = BasisFunctions(wfs.gd,
                         [setup.basis_functions_J for setup in wfs.setups],
                         wfs.kd, cut=True)
    bfs.set_positions(wfs.spos_ac)
    return bfs


def get_lcao_projections_HSP(calc, bfs=None, spin=0, projectionsonly=True):
    """Some title.

    if projectionsonly is True, return the projections::

      V_qnM = <psi_qn | Phi_qM>

    else, also return the Hamiltonian, overlap, and projector overlaps::

      H_qMM  = <Phi_qM| H |Phi_qM'>
      S_qMM  = <Phi_qM|Phi_qM'>
      P_aqMi = <pt^a_qi|Phi_qM>
    """
    if calc.wfs.kd.comm.size != 1:
        raise NotImplementedError('Parallelization over spin/kpt not '
                                  'implemented yet.')
    comm = calc.wfs.gd.comm
    nq = len(calc.wfs.kd.ibzk_qc)
    nao = calc.wfs.setups.nao
    dtype = calc.wfs.dtype
    if bfs is None:
        bfs = get_bfs(calc)
    tci = TwoCenterIntegrals(calc.wfs.gd.cell_cv,
                             calc.wfs.gd.pbc_c,
                             calc.wfs.setups,
                             calc.wfs.kd.ibzk_qc,
                             calc.wfs.kd.gamma)

    # Calculate projector overlaps, and (lower triangle of-) S and T matrices
    S_qMM = np.zeros((nq, nao, nao), dtype)
    T_qMM = np.zeros((nq, nao, nao), dtype)
    P_aqMi = {}

    for a in bfs.my_atom_indices:
        ni = calc.wfs.setups[a].ni
        P_aqMi[a] = np.zeros((nq, nao, ni), dtype)
    tci.calculate(calc.spos_ac, S_qMM, T_qMM, P_aqMi)

    from gpaw.lcao.atomic_correction import DenseAtomicCorrection
    dS_aii = {a: calc.wfs.setups[a].dO_ii for a in bfs.my_atom_indices}
    atomic_correction = DenseAtomicCorrection(P_aqMi, dS_aii, 0, nao)
    atomic_correction.add_overlap_correction(S_qMM)
    # calc.wfs.P_aqMi = P_aqMi
    # for q, S_MM in enumerate(S_qMM):
    #     atomic_correction.calculate_manual(wfs, q,
    # atomic_correction.add_overlap_correction(calc.wfs, S_qMM)
    calc.wfs.gd.comm.sum(S_qMM)
    calc.wfs.gd.comm.sum(T_qMM)

    # Calculate projections
    V_qnM = np.zeros((nq, calc.wfs.bd.nbands, nao), dtype)
    for kpt in calc.wfs.kpt_u:
        if kpt.s != spin:
            continue
        V_nM = V_qnM[kpt.q]
        # bfs.integrate2(kpt.psit_nG[:], V_nM, kpt.q) # all bands to save time
        for n, V_M in enumerate(V_nM):  # band-by-band to save memory
            bfs.integrate2(kpt.psit_nG[n][:], V_M, kpt.q)
        for a, P_ni in kpt.projections.items():
            dS_ii = calc.wfs.setups[a].dO_ii
            P_Mi = P_aqMi[a][kpt.q]
            V_nM += np.dot(P_ni, np.inner(dS_ii, P_Mi).conj())
    comm.sum(V_qnM)
    if projectionsonly:
        return V_qnM

    # Determine potential matrix
    vt_G = calc.hamiltonian.vt_sG[spin]
    H_qMM = np.zeros((nq, nao, nao), dtype)
    for q, H_MM in enumerate(H_qMM):
        bfs.calculate_potential_matrix(vt_G, H_MM, q)

    # Make Hamiltonian as sum of kinetic (T) and potential (V) matrices
    # and add atomic corrections
    for a, P_qMi in P_aqMi.items():
        dH_ii = unpack_hermitian(calc.hamiltonian.dH_asp[a][spin])
        for P_Mi, H_MM in zip(P_qMi, H_qMM):
            H_MM += np.dot(P_Mi, np.inner(dH_ii, P_Mi).conj())
    comm.sum(H_qMM)
    H_qMM += T_qMM

    # Fill in the upper triangles of H and S
    for H_MM, S_MM in zip(H_qMM, S_qMM):
        tri2full(H_MM)
        tri2full(S_MM)
    H_qMM *= Hartree

    return V_qnM, H_qMM, S_qMM, P_aqMi


def convert_projection_data(symbols, V_qnM, H_qMM, S_qMM, P_aqMi,
                            originaltype='dzp', newtype='sz'):
    Nq = len(H_qMM)
    mask = basis_subset2(symbols, originaltype, newtype)
    mask2 = np.outer(mask, mask)
    NMnew = len(mask.nonzero()[0])
    V_qnM = V_qnM[..., mask]
    H_qMM = H_qMM[:, mask2].reshape(Nq, NMnew, NMnew)
    S_qMM = S_qMM[:, mask2].reshape(Nq, NMnew, NMnew)
    P2_aqMi = {}
    for a, P_qMi in P_aqMi.items():
        P2_aqMi[a] = P_qMi[:, mask, :]
    return V_qnM, H_qMM, S_qMM, P2_aqMi


def get_phs(calc, s=0):
    return get_lcao_projections_HSP(calc, bfs=None,
                                    spin=s, projectionsonly=False)


class ProjectedWannierFunctions:
    def __init__(self, projections, h_lcao, s_lcao, eigenvalues, kpoints,
                 L_k=None, M_k=None, N=None, fixedenergy=None):
        """projections[n,i] = <psi_n|f_i>
           h_lcao[i1, i2] = <f_i1|h|f_i2>
           s_lcao[[i1, i2] = <f_i1|f_i2>
           eps_n: Exact eigenvalues
           L: Number of extra degrees of freedom
           M: Number of states to exactly span
           N: Total number of bands in the calculation

           Methods:
           -- get_hamiltonian_and_overlap_matrix --
           will return the hamiltonian and identity operator
           in the projected wannier function basis.
           The following steps are performed:

           1) calculate_edf       -> self.b_il
           2) calculate_rotations -> self.Uo_mi and self.Uu_li
           3) calculate_overlaps  -> self.S_ii
           4) calculate_hamiltonian_matrix -> self.H_ii

           -- get_eigenvalues --
           gives the eigenvalues of of the hamiltonian in the
           projected wannier function basis.

           -- indices --
           i localized function index
           n eigenstate index
           l edf index
           m fixed eigenstate index
           k k-point index
           """

        self.eps_kn = eigenvalues
        self.ibzk_kc = kpoints
        self.nk = len(self.ibzk_kc)
        self.V_kni = projections  # <psi_n1|f_i1>
        self.dtype = self.V_kni.dtype
        self.Nw = self.V_kni.shape[2]
        self.s_lcao_kii = s_lcao  # F_ii[i1,i2] = <f_i1|f_i2>
        self.h_lcao_kii = h_lcao

        if N is None:
            N = self.V_kni.shape[1]
        self.N = N

        if fixedenergy is None:
            raise NotImplementedError(
                'Only fixedenergy is implemented for now')
        else:
            self.fixedenergy = fixedenergy
            self.M_k = [sum(eps_n <= fixedenergy) for eps_n in self.eps_kn]
            self.L_k = [self.Nw - M for M in self.M_k]
            print("fixedenergy =", self.fixedenergy)

        print('N =', self.N)
        print('skpt_kc = ')
        print(self.ibzk_kc)
        print('M_k =', self.M_k)
        print('L_k =', self.L_k)
        print('Nw =', self.Nw)

    def get_hamiltonian_and_overlap_matrix(self, useibl=True):
        self.calculate_edf(useibl=useibl)
        self.calculate_rotations()
        self.calculate_overlaps()
        self.calculate_hamiltonian_matrix(useibl=useibl)
        return self.H_kii, self.S_kii

    def calculate_edf(self, useibl=True):
        """Calculate the coefficients b_il in the expansion of the EDF.

        ``|phi_l> = sum_i b_il |f^u_i>``, in terms of ``|f^u_i> = P^u|f_i>``.

        To use the infinite band limit set useibl=True.
        N is the total number of bands to use.
        """

        for k, L in enumerate(self.L_k):
            if L == 0:
                assert L != 0, 'L_k=0 for k=%i. Not implemented' % k

        self.Vo_kni = [V_ni[:M] for V_ni, M in zip(self.V_kni, self.M_k)]

        self.Fo_kii = np.asarray([np.dot(dagger(Vo_ni), Vo_ni)
                                  for Vo_ni in self.Vo_kni])

        if useibl:
            self.Fu_kii = self.s_lcao_kii - self.Fo_kii
        else:
            self.Vu_kni = [V_ni[M:self.N]
                           for V_ni, M in zip(self.V_kni, self.M_k)]
            self.Fu_kii = np.asarray([np.dot(dagger(Vu_ni), Vu_ni)
                                     for Vu_ni in self.Vu_kni])
        self.b_kil = []
        for Fu_ii, L in zip(self.Fu_kii, self.L_k):
            b_i, b_ii = la.eigh(Fu_ii)
            ls = b_i.real.argsort()[-L:]
            b_il = b_ii[:, ls]  # pick out the eigenvec with largest eigenvals.
            normalize2(b_il, Fu_ii)  # normalize the EDF: <phi_l|phi_l> = 1
            self.b_kil.append(b_il)

    def calculate_rotations(self):
        Uo_kni = [Vo_ni.copy() for Vo_ni in self.Vo_kni]
        Uu_kli = [np.dot(dagger(b_il), Fu_ii)
                  for b_il, Fu_ii in zip(self.b_kil, self.Fu_kii)]
        # Normalize such that <omega_i|omega_i> = <f_i|f_i>
        for Uo_ni, Uu_li, s_ii in zip(Uo_kni, Uu_kli, self.s_lcao_kii):
            normalize(Uo_ni, Uu_li, s_ii.diagonal())
        self.Uo_kni = Uo_kni
        self.Uu_kli = Uu_kli

    def calculate_overlaps(self):
        Wo_kii = [np.dot(dagger(Uo_ni), Uo_ni) for Uo_ni in self.Uo_kni]
        Wu_kii = [np.dot(dagger(Uu_li), Uu_li) for Uu_li in self.Uu_kli]
        # Wu_kii = [dots(dagger(Uu_li), dagger(b_il), Fu_ii, b_il, Uu_li)
        # for Uu_li, b_il, Fu_ii in zip(self.Uu_kli, self.b_kil, self.Fu_kii)]
        self.S_kii = np.asarray(Wo_kii) + np.asarray(Wu_kii)

    def get_condition_number(self):
        return np.asarray([condition_number(S) for S in self.S_kii])

    def calculate_hamiltonian_matrix(self, useibl=True):
        """Calculate H_kij = H^o_i(k)j(k) + H^u_i(k)j(k)
           i(k): Bloch sum of omega_i
        """

        epso_kn = [eps_n[:M] for eps_n, M in zip(self.eps_kn, self.M_k)]
        self.Ho_kii = np.asarray([np.dot(dagger(Uo_ni) * epso_n, Uo_ni)
                                  for Uo_ni, epso_n in zip(self.Uo_kni,
                                                           epso_kn)])

        if self.h_lcao_kii is not None and useibl:
            print("Using h_lcao and infinite band limit")
            Huf_kii = [h_lcao_ii - np.dot(dagger(Vo_ni) * epso_n, Vo_ni)
                       for h_lcao_ii, Vo_ni, epso_n in zip(self.h_lcao_kii,
                                                           self.Vo_kni,
                                                           epso_kn)]
            self.Huf_kii = np.asarray(Huf_kii)
        else:
            print("Using finite band limit (not using h_lcao)")
            epsu_kn = [eps_n[M:self.N]
                       for eps_n, M in zip(self.eps_kn, self.M_k)]
            Huf_kii = [np.dot(dagger(Vu_ni) * epsu_n, Vu_ni)
                       for Vu_ni, epsu_n in zip(self.Vu_kni, epsu_kn)]
            self.Huf_kii = np.asarray(Huf_kii)

        Hu_kii = [dots(dagger(Uu_li), dagger(b_il), Huf_ii, b_il, Uu_li)
                  for Uu_li, b_il, Huf_ii in zip(self.Uu_kli, self.b_kil,
                                                 self.Huf_kii)]
        self.Hu_kii = np.asarray(Hu_kii)
        self.H_kii = self.Ho_kii + self.Hu_kii

    def get_eigenvalues(self):
        return np.asarray([eigvals(H, S)
                           for H, S in zip(self.H_kii, self.S_kii)])

    def get_lcao_eigenvalues(self):
        return np.asarray([eigvals(H, S)
                           for H, S in zip(self.h_lcao_kii, self.s_lcao_kii)])

    def get_norm_of_projection(self):
        Sinv_kii = np.asarray([la.inv(S_ii) for S_ii in self.S_kii])

        normo_kn = np.asarray([dots(Uo_ni, Sinv_ii, dagger(Uo_ni)).diagonal()
                               for Uo_ni, Sinv_ii in
                               zip(self.Uo_kni, Sinv_kii)])

        Vu_kni = np.asarray([V_ni[M:self.N]
                             for V_ni, M in zip(self.V_kni, self.M_k)])

        Pu_kni = [dots(Vu_ni, b_il, Uu_li)
                  for Vu_ni, b_il, Uu_li in
                  zip(Vu_kni, self.b_kil, self.Uu_kli)]

        normu_kn = np.asarray([dots(Pu_ni, Sinv_ii, dagger(Pu_ni)).diagonal()
                               for Pu_ni, Sinv_ii in zip(Pu_kni, Sinv_kii)])

        return np.hstack((normo_kn, normu_kn))

    def calculate_functions(self, calc, basis, k=0):
        from gpaw.io import FileReference
        psit_nG = calc.wfs.kpt_u[k].psit_nG
        Uo_ni = self.Uo_kni[k]
        tarinstance = isinstance(psit_nG, FileReference)
        if tarinstance:
            psit_nG = np.asarray([psit_nG[i] for i in range(self.M_k[k])])

        basis_functions = get_bfs(calc)
        b_il, Uu_li, Vo_ni = self.b_kil[k], self.Uu_kli[k], self.Vo_kni[k]
        a_iG = -np.tensordot(Vo_ni, psit_nG, axes=[0, 0])  # a_iG = -fo_iG
        # self.fo_iG = -a_iG.copy()#f projected onto the occupied subspace
        C_iM = np.identity(self.Nw, dtype=self.dtype)
        basis_functions.lcao_to_grid(C_iM, a_iG, q=-1)  # a_iG=fu_iG=f_iG-fo_iG
        # self.f_iG = self.fo_iG + a_iG # check
        a_iG = np.tensordot(b_il, a_iG, axes=[0, 0])  # a_iG = EDF
        a_iG = np.tensordot(Uu_li, a_iG, axes=[0, 0])  # a_iG = wu_iG
        a_iG += np.tensordot(Uo_ni, psit_nG, axes=[0, 0])  # ai_G = wu_iG+wo_iG
        self.w_iG = a_iG

    def get_mlwf_initial_guess(self):
        """calculate initial guess for maximally localized
        wannier functions. Does not work for the infinite band limit.
        cu_nl: rotation coefficents of unoccupied states
        U_ii: rotation matrix of eigenstates and edf.
        """
        Vu_ni = self.Vu_ni[self.M: self.N]
        cu_nl = np.dot(Vu_ni, self.b_il)
        U_ii = np.vstack((self.Uo_ni, self.Uu_li))
        lowdin(U_ii)
        return U_ii, cu_nl
