import numpy as np
from typing import NamedTuple

from ase.units import Hartree, Bohr

from ase.io.ulm import Reader
from gpaw.io import Writer
from gpaw.external import ConstantElectricField
from gpaw.kpoint import KPoint
from gpaw.lcaotddft.hamiltonian import KickHamiltonian
from gpaw.lcaotddft.utilities import collect_MM
from gpaw.lcaotddft.utilities import distribute_nM
from gpaw.lcaotddft.utilities import read_uMM
from gpaw.lcaotddft.utilities import write_uMM
from gpaw.lcaotddft.utilities import read_uX, write_uX
from gpaw.utilities.scalapack import \
    pblas_simple_gemm, pblas_simple_hemm, scalapack_tri2full
from gpaw.utilities.tools import tri2full


def gauss_ij(energy_i, energy_j, sigma):
    denergy_ij = energy_i[:, np.newaxis] - energy_j[np.newaxis, :]
    norm = 1.0 / (sigma * np.sqrt(2 * np.pi))
    return norm * np.exp(-0.5 * denergy_ij**2 / sigma**2)


def get_bfs_maps(calc):
    # Construct maps
    # a_M: M -> atom index a
    # l_M: M -> angular momentum l
    a_M = []
    l_M = []
    M = 0
    for a, sphere in enumerate(calc.wfs.basis_functions.sphere_a):
        for j, spline in enumerate(sphere.spline_j):
            l = spline.get_angular_momentum_number()
            for _ in range(2 * l + 1):
                a_M.append(a)
                l_M.append(l)
                M += 1
    a_M = np.array(a_M)
    l_M = np.array(l_M)
    return a_M, l_M


class DummyKsl(NamedTuple):
    using_blacs: bool = False


class KohnShamDecomposition:
    version = 1
    ulmtag = 'KSD'
    readwrite_attrs = ['fermilevel', 'only_ia', 'w_p', 'f_p', 'ia_p',
                       'P_p', 'dm_vp', 'a_M', 'l_M']

    def __init__(self, paw=None, filename=None):
        self.filename = filename
        self.has_initialized = False
        self.reader = None
        if paw is not None:
            self.world = paw.world
            self.log = paw.log
            self.ksl = paw.wfs.ksl
            self.kd = paw.wfs.kd
            self.bd = paw.wfs.bd
            self.kpt_u = paw.wfs.kpt_u
            self.density = paw.density
            self.comm = paw.comms['K']

            if len(paw.wfs.kpt_u) > 1:
                raise RuntimeError('K-points are not fully supported')

        if filename is not None:
            self.read(filename)

            if paw is not None:
                return

            # Create a dummy KohnShamLayouts object and one Gamma k-point
            # This is necessary to read attributes
            ns, nk = self.reader.eig_un.shape[:2]
            assert (ns, nk) == (1, 1), 'Spins and K-points not implemented'
            self.kpt_u = [KPoint(weightk=1, weight=1, s=0, k=0, q=0)]
            self.ksl = DummyKsl()

    def initialize(self, paw, min_occdiff=1e-3, only_ia=True):
        if self.has_initialized:
            return
        paw.initialize_positions()
        # paw.set_positions()

        assert self.bd.nbands == self.ksl.nao
        self.only_ia = only_ia

        if not self.ksl.using_blacs and self.bd.comm.size > 1:
            raise RuntimeError('Band parallelization without scalapack '
                               'is not supported')

        if self.kd.gamma:
            self.C0_dtype = float
        else:
            self.C0_dtype = complex

        # Take quantities
        self.fermilevel = paw.wfs.fermi_level
        self.S_uMM = []
        self.C0_unM = []
        self.eig_un = []
        self.occ_un = []
        for kpt in paw.wfs.kpt_u:
            S_MM = kpt.S_MM
            assert np.max(np.absolute(S_MM.imag)) == 0.0
            S_MM = np.ascontiguousarray(S_MM.real)
            if self.ksl.using_blacs:
                scalapack_tri2full(self.ksl.mmdescriptor, S_MM)
            self.S_uMM.append(S_MM)

            C_nM = kpt.C_nM
            if self.C0_dtype == float:
                assert np.max(np.absolute(C_nM.imag)) == 0.0
                C_nM = np.ascontiguousarray(C_nM.real)
            C_nM = distribute_nM(self.ksl, C_nM)
            self.C0_unM.append(C_nM)

            eig_n = paw.wfs.collect_eigenvalues(kpt.k, kpt.s)
            occ_n = paw.wfs.collect_occupations(kpt.k, kpt.s)
            self.eig_un.append(eig_n)
            self.occ_un.append(occ_n)

        self.a_M, self.l_M = get_bfs_maps(paw)
        self.atoms = paw.atoms

        # TODO: do the rest of the function with K-points

        # Construct p = (i, a) pairs
        u = 0
        eig_n = self.eig_un[u]
        occ_n = self.occ_un[u]
        C0_nM = self.C0_unM[u]

        if self.comm.rank == 0:
            Nn = self.bd.nbands

            f_p = []
            w_p = []
            i_p = []
            a_p = []
            ia_p = []
            i0 = 0
            for i in range(i0, Nn):
                if only_ia:
                    a0 = i + 1
                else:
                    a0 = 0
                for a in range(a0, Nn):
                    f = occ_n[i] - occ_n[a]
                    if only_ia and f < min_occdiff:
                        continue
                    w = eig_n[a] - eig_n[i]
                    f_p.append(f)
                    w_p.append(w)
                    i_p.append(i)
                    a_p.append(a)
                    ia_p.append((i, a))
            f_p = np.array(f_p)
            w_p = np.array(w_p)
            i_p = np.array(i_p, dtype=int)
            a_p = np.array(a_p, dtype=int)
            ia_p = np.array(ia_p, dtype=int)

            # Sort according to energy difference
            p_s = np.argsort(w_p)
            f_p = f_p[p_s]
            w_p = w_p[p_s]
            i_p = i_p[p_s]
            a_p = a_p[p_s]
            ia_p = ia_p[p_s]

            Np = len(f_p)
            P_p = []
            for p in range(Np):
                P = np.ravel_multi_index(ia_p[p], (Nn, Nn))
                P_p.append(P)
            P_p = np.array(P_p)

            dm_vp = np.empty((3, Np), dtype=float)

        for v in range(3):
            direction = np.zeros(3, dtype=float)
            direction[v] = 1.0
            cef = ConstantElectricField(Hartree / Bohr, direction)
            kick_hamiltonian = KickHamiltonian(paw.hamiltonian, paw.density,
                                               cef)
            dm_MM = paw.wfs.eigensolver.calculate_hamiltonian_matrix(
                kick_hamiltonian, paw.wfs, paw.wfs.kpt_u[u],
                add_kinetic=False, root=-1)

            if self.ksl.using_blacs:
                tmp_nM = self.ksl.mmdescriptor.zeros(dtype=C0_nM.dtype)
                pblas_simple_hemm(self.ksl.mmdescriptor,
                                  self.ksl.mmdescriptor,
                                  self.ksl.mmdescriptor,
                                  dm_MM, C0_nM.conj(), tmp_nM,
                                  side='R', uplo='L')
                dm_nn = self.ksl.mmdescriptor.zeros(dtype=C0_nM.dtype)
                pblas_simple_gemm(self.ksl.mmdescriptor,
                                  self.ksl.mmdescriptor,
                                  self.ksl.mmdescriptor,
                                  tmp_nM, C0_nM, dm_nn, transb='T')
            else:
                tri2full(dm_MM)
                dm_nn = np.dot(C0_nM.conj(), np.dot(dm_MM, C0_nM.T))

            dm_nn = collect_MM(self.ksl, dm_nn)
            if self.comm.rank == 0:
                dm_P = dm_nn.ravel()
                dm_p = dm_P[P_p]
                dm_vp[v] = dm_p

        if self.comm.rank == 0:
            self.w_p = w_p
            self.f_p = f_p
            self.ia_p = ia_p
            self.P_p = P_p
            self.dm_vp = dm_vp

        self.has_initialized = True

    def write(self, filename):
        from ase.io.trajectory import write_atoms

        self.log(f'{self.__class__.__name__}: Writing to {filename}')
        writer = Writer(filename, self.world, mode='w',
                        tag=self.__class__.ulmtag)
        writer.write(version=self.__class__.version)

        write_atoms(writer.child('atoms'), self.atoms)

        writer.write(ha=Hartree)
        write_uMM(self.kd, self.ksl, writer, 'S_uMM', self.S_uMM)
        write_uMM(self.kd, self.ksl, writer, 'C0_unM', self.C0_unM)
        write_uX(self.kd, self.ksl.block_comm, writer, 'eig_un', self.eig_un)
        write_uX(self.kd, self.ksl.block_comm, writer, 'occ_un', self.occ_un)

        if self.comm.rank == 0:
            for arg in self.readwrite_attrs:
                writer.write(arg, getattr(self, arg))

        writer.close()

    def read(self, filename):
        self.reader = Reader(filename)
        tag = self.reader.get_tag()
        if tag != self.__class__.ulmtag:
            raise RuntimeError('Unknown tag %s' % tag)
        self.version = self.reader.version

        # Do lazy reading in __getattr__ only if/when
        # the variables are required
        self.has_initialized = True

    def __getattr__(self, attr):
        if attr in ['S_uMM', 'C0_unM']:
            val = read_uMM(self.kpt_u, self.ksl, self.reader, attr)
            setattr(self, attr, val)
            return val
        if attr in ['eig_un', 'occ_un']:
            val = read_uX(self.kpt_u, self.reader, attr)
            setattr(self, attr, val)
            return val
        if attr in ['C0S_unM']:
            C0S_unM = []
            for u, kpt in enumerate(self.kpt_u):
                C0_nM = self.C0_unM[u]
                S_MM = self.S_uMM[u]
                if self.ksl.using_blacs:
                    C0S_nM = self.ksl.mmdescriptor.zeros(dtype=C0_nM.dtype)
                    pblas_simple_hemm(self.ksl.mmdescriptor,
                                      self.ksl.mmdescriptor,
                                      self.ksl.mmdescriptor,
                                      S_MM, C0_nM, C0S_nM,
                                      side='R', uplo='L')
                else:
                    C0S_nM = np.dot(C0_nM, S_MM)
                C0S_unM.append(C0S_nM)
            setattr(self, attr, C0S_unM)
            return C0S_unM
        if attr in ['weight_Mn']:
            assert self.world.size == 1
            C2_nM = np.absolute(self.C0_unM[0])**2
            val = C2_nM.T / np.sum(C2_nM, axis=1)
            setattr(self, attr, val)
            return val

        try:
            val = getattr(self.reader, attr)
            if attr == 'atoms':
                from ase.io.trajectory import read_atoms
                val = read_atoms(val)
            setattr(self, attr, val)
            return val
        except (KeyError, AttributeError):
            pass

        raise AttributeError('Attribute %s not defined in version %s' %
                             (repr(attr), repr(self.version)))

    def distribute(self, comm):
        self.comm = comm
        N = comm.size
        self.Np = len(self.P_p)
        self.Nq = int(np.ceil(self.Np / float(N)))
        self.NQ = self.Nq * N
        self.w_q = self.distribute_p(self.w_p)
        self.f_q = self.distribute_p(self.f_p)
        self.dm_vq = self.distribute_xp(self.dm_vp)

    def distribute_p(self, a_p, a_q=None, root=0):
        if a_q is None:
            a_q = np.zeros(self.Nq, dtype=a_p.dtype)
        if self.comm.rank == root:
            a_Q = np.append(a_p, np.zeros(self.NQ - self.Np, dtype=a_p.dtype))
        else:
            a_Q = None
        self.comm.scatter(a_Q, a_q, root)
        return a_q

    def collect_q(self, a_q, root=0):
        if self.comm.rank == root:
            a_Q = np.zeros(self.NQ, dtype=a_q.dtype)
        else:
            a_Q = None
        self.comm.gather(a_q, root, a_Q)
        if self.comm.rank == root:
            a_p = a_Q[:self.Np]
        else:
            a_p = None
        return a_p

    def distribute_xp(self, a_xp):
        Nx = a_xp.shape[0]
        a_xq = np.zeros((Nx, self.Nq), dtype=a_xp.dtype)
        for x in range(Nx):
            self.distribute_p(a_xp[x], a_xq[x])
        return a_xq

    def transform(self, rho_uMM, broadcast=False):
        assert len(rho_uMM) == 1, 'K-points not implemented'
        u = 0
        rho_MM = np.ascontiguousarray(rho_uMM[u])
        C0S_nM = self.C0S_unM[u].astype(rho_MM.dtype, copy=True)
        # KS decomposition
        if self.ksl.using_blacs:
            tmp_nM = self.ksl.mmdescriptor.zeros(dtype=rho_MM.dtype)
            pblas_simple_gemm(self.ksl.mmdescriptor,
                              self.ksl.mmdescriptor,
                              self.ksl.mmdescriptor,
                              C0S_nM, rho_MM, tmp_nM)
            rho_nn = self.ksl.mmdescriptor.zeros(dtype=rho_MM.dtype)
            pblas_simple_gemm(self.ksl.mmdescriptor,
                              self.ksl.mmdescriptor,
                              self.ksl.mmdescriptor,
                              tmp_nM, C0S_nM, rho_nn, transb='C')
        else:
            rho_nn = np.dot(np.dot(C0S_nM, rho_MM), C0S_nM.T.conj())

        rho_nn = collect_MM(self.ksl, rho_nn)
        if self.comm.rank == 0:
            rho_P = rho_nn.ravel()
            # Remove de-excitation terms
            rho_p = rho_P[self.P_p]
            if self.only_ia:
                rho_p *= 2
        else:
            rho_p = None

        if broadcast:
            if self.comm.rank != 0:
                rho_p = np.zeros_like(self.P_p, dtype=rho_MM.dtype)
            self.comm.broadcast(rho_p, 0)
        rho_up = [rho_p]
        return rho_up

    def ialims(self):
        i_p = self.ia_p[:, 0]
        a_p = self.ia_p[:, 1]
        imin = np.min(i_p)
        imax = np.max(i_p)
        amin = np.min(a_p)
        amax = np.max(a_p)
        return imin, imax, amin, amax

    def M_p_to_M_ia(self, M_p):
        return self.M_ia_from_M_p(M_p)

    def M_ia_from_M_p(self, M_p):
        imin, imax, amin, amax = self.ialims()
        M_ia = np.zeros((imax - imin + 1, amax - amin + 1), dtype=M_p.dtype)
        for M, (i, a) in zip(M_p, self.ia_p):
            M_ia[i - imin, a - amin] = M
        return M_ia

    def plot_matrix(self, M_p):
        import matplotlib.pyplot as plt
        M_ia = self.M_ia_from_M_p(M_p)
        plt.imshow(M_ia, interpolation='none')
        plt.xlabel('a')
        plt.ylabel('i')

    def get_dipole_moment_contributions(self, rho_up):
        assert len(rho_up) == 1, 'K-points not implemented'
        u = 0
        rho_p = rho_up[u]
        dmrho_vp = - self.dm_vp * rho_p
        return dmrho_vp

    def get_dipole_moment(self, rho_up):
        assert len(rho_up) == 1, 'K-points not implemented'
        u = 0
        rho_p = rho_up[u]
        dm_v = - np.dot(self.dm_vp, rho_p)
        return dm_v

    def get_density(self, wfs, rho_up, density='comp'):
        from gpaw.lcaotddft.densitymatrix import get_density

        if self.ksl.using_blacs:
            raise NotImplementedError('Scalapack is not supported')

        density_type = density
        assert len(rho_up) == 1, 'K-points not implemented'
        u = 0
        rho_p = rho_up[u]
        C0_nM = self.C0_unM[u]

        rho_ia = self.M_ia_from_M_p(rho_p)
        imin, imax, amin, amax = self.ialims()
        C0_iM = C0_nM[imin:(imax + 1)]
        C0_aM = C0_nM[amin:(amax + 1)]

        rho_MM = np.dot(C0_iM.T, np.dot(rho_ia, C0_aM.conj()))
        rho_MM = 0.5 * (rho_MM + rho_MM.T)

        return get_density(rho_MM, wfs, self.density, density_type, u)

    def get_contributions_table(self, weight_p, minweight=0.01,
                                zero_fermilevel=True):
        assert weight_p.dtype == float
        u = 0  # TODO

        absweight_p = np.absolute(weight_p)
        tot_weight = weight_p.sum()
        propweight_p = weight_p / tot_weight * 100
        tot_propweight = propweight_p.sum()
        rest_weight = tot_weight
        rest_propweight = tot_propweight
        eig_n = self.eig_un[u].copy()
        if zero_fermilevel:
            eig_n -= self.fermilevel

        txt = ''
        txt += ('# %6s %4s(%8s)    %4s(%8s)  %12s %14s %8s\n' %
                ('p', 'i', 'eV', 'a', 'eV', 'Ediff (eV)', 'weight', '%'))
        p_s = np.argsort(absweight_p)[::-1]
        for s, p in enumerate(p_s):
            i, a = self.ia_p[p]
            if absweight_p[p] < minweight:
                break
            txt += ('  %6s %4d(%8.3f) -> %4d(%8.3f): %12.4f %14.4f %8.1f\n' %
                    (p, i, eig_n[i] * Hartree, a, eig_n[a] * Hartree,
                     self.w_p[p] * Hartree, weight_p[p], propweight_p[p]))
            rest_weight -= weight_p[p]
            rest_propweight -= propweight_p[p]
        txt += ('  %39s: %12s %+14.4f %8.1f\n' %
                ('rest', '', rest_weight, rest_propweight))
        txt += ('  %39s: %12s %+14.4f %8.1f\n' %
                ('total', '', tot_weight, tot_propweight))
        return txt

    def plot_TCM(self, weight_p, energy_o, energy_u, sigma,
                 zero_fermilevel=True, vmax='80%'):
        from gpaw.lcaotddft.tcm import TCMPlotter
        plotter = TCMPlotter(self, energy_o, energy_u, sigma, zero_fermilevel)
        ax_tcm = plotter.plot_TCM(weight_p, vmax)
        ax_occ_dos, ax_unocc_dos = plotter.plot_DOS()
        return ax_tcm, ax_occ_dos, ax_unocc_dos

    def get_TCM(self, weight_p, eig_n, energy_o, energy_u, sigma):
        flt_p = self.filter_by_x_ia(eig_n, energy_o, energy_u, 8 * sigma)
        weight_f = weight_p[flt_p]
        G_fo = gauss_ij(eig_n[self.ia_p[flt_p, 0]], energy_o, sigma)
        G_fu = gauss_ij(eig_n[self.ia_p[flt_p, 1]], energy_u, sigma)
        tcm_ou = np.dot(G_fo.T * weight_f, G_fu)
        return tcm_ou

    def get_DOS(self, eig_n, energy_o, energy_u, sigma):
        return self.get_weighted_DOS(1, eig_n, energy_o, energy_u, sigma)

    def get_weighted_DOS(self, weight_n, eig_n, energy_o, energy_u, sigma):
        if not isinstance(weight_n, np.ndarray):
            # Assume float
            weight_n = weight_n * np.ones_like(eig_n)
        G_on = gauss_ij(energy_o, eig_n, sigma)
        G_un = gauss_ij(energy_u, eig_n, sigma)
        dos_o = np.dot(G_on, weight_n)
        dos_u = np.dot(G_un, weight_n)
        return dos_o, dos_u

    def get_weight_n_by_l(self, l):
        if isinstance(l, int):
            weight_n = np.sum(self.weight_Mn[self.l_M == l], axis=0)
        else:
            weight_n = np.sum([self.get_weight_n_by_l(l_) for l_ in l],
                              axis=0)
        return weight_n

    def get_weight_n_by_a(self, a):
        if isinstance(a, int):
            weight_n = np.sum(self.weight_Mn[self.a_M == a], axis=0)
        else:
            weight_n = np.sum([self.get_weight_n_by_a(a_) for a_ in a],
                              axis=0)
        return weight_n

    def get_distribution_i(self, weight_p, energy_e, sigma,
                           zero_fermilevel=True):
        eig_n, fermilevel = self.get_eig_n(zero_fermilevel)
        flt_p = self.filter_by_x_i(eig_n, energy_e, 8 * sigma)
        weight_f = weight_p[flt_p]
        G_fe = gauss_ij(eig_n[self.ia_p[flt_p, 0]], energy_e, sigma)
        dist_e = np.dot(G_fe.T, weight_f)
        return dist_e

    def get_distribution_a(self, weight_p, energy_e, sigma,
                           zero_fermilevel=True):
        eig_n, fermilevel = self.get_eig_n(zero_fermilevel)
        flt_p = self.filter_by_x_a(eig_n, energy_e, 8 * sigma)
        weight_f = weight_p[flt_p]
        G_fe = gauss_ij(eig_n[self.ia_p[flt_p, 1]], energy_e, sigma)
        dist_e = np.dot(G_fe.T, weight_f)
        return dist_e

    def get_distribution_ia(self, weight_p, energy_o, energy_u, sigma,
                            zero_fermilevel=True):
        """
        Filter both i and a spaces as in TCM.

        """
        eig_n, fermilevel = self.get_eig_n(zero_fermilevel)
        flt_p = self.filter_by_x_ia(eig_n, energy_o, energy_u, 8 * sigma)
        weight_f = weight_p[flt_p]
        G_fo = gauss_ij(eig_n[self.ia_p[flt_p, 0]], energy_o, sigma)
        dist_o = np.dot(G_fo.T, weight_f)
        G_fu = gauss_ij(eig_n[self.ia_p[flt_p, 1]], energy_u, sigma)
        dist_u = np.dot(G_fu.T, weight_f)
        return dist_o, dist_u

    def get_distribution(self, weight_p, energy_e, sigma):
        w_p = self.w_p * Hartree
        flt_p = self.filter_by_x_p(w_p, energy_e, 8 * sigma)
        weight_f = weight_p[flt_p]
        G_fe = gauss_ij(w_p[flt_p], energy_e, sigma)
        dist_e = np.dot(G_fe.T, weight_f)
        return dist_e

    def get_eig_n(self, zero_fermilevel=True):
        u = 0  # TODO
        eig_n = self.eig_un[u].copy()
        if zero_fermilevel:
            eig_n -= self.fermilevel
            fermilevel = 0.0
        else:
            fermilevel = self.fermilevel
        eig_n *= Hartree
        fermilevel *= Hartree
        return eig_n, fermilevel

    def filter_by_x_p(self, x_p, energy_e, buf):
        flt_p = np.logical_and((energy_e[0] - buf) <= x_p,
                               x_p <= (energy_e[-1] + buf))
        return flt_p

    def filter_by_x_i(self, x_n, energy_e, buf):
        return self.filter_by_x_p(x_n[self.ia_p[:, 0]], energy_e, buf)

    def filter_by_x_a(self, x_n, energy_e, buf):
        return self.filter_by_x_p(x_n[self.ia_p[:, 1]], energy_e, buf)

    def filter_by_x_ia(self, x_n, energy_o, energy_u, buf):
        flti_p = self.filter_by_x_i(x_n, energy_o, buf)
        flta_p = self.filter_by_x_a(x_n, energy_u, buf)
        flt_p = np.logical_and(flti_p, flta_p)
        return flt_p

    def __del__(self):
        if self.reader is not None:
            self.reader.close()
