from math import pi, sqrt
import numpy as np
from ase.parallel import paropen
from gpaw.utilities import pack_density
from gpaw.analyse.wignerseitz import wignerseitz
from gpaw.setup_data import SetupData
from gpaw.gauss import Gauss
from gpaw.io.fmf import FMF
from gpaw.utilities.blas import gemmdot


def print_projectors(setup):
    """Print information on the projectors of input nucleus object.

    If nucleus is a string, treat this as an element name.
    """
    if isinstance(setup, str):
        setup = SetupData(setup, 'LDA', 'paw')
        n_j = setup.n_j
        l_j = setup.l_j
    else:
        n_j = setup.n_j
        l_j = setup.l_j

    angular = [['1'],
               ['y', 'z', 'x'],
               ['xy', 'yz', '3z^2-r^2', 'xz', 'x^2-y^2'],
               ['3x^2y-y^3', 'xyz', '5yz^2-yr^2', '5z^3-3zr^2',
                '5xz^2-xr^2', 'x^2z-y^2z', 'x^3-3xy^2'],
               ]
    print(' i n l m')
    print('--------')
    i = 0
    for n, l in zip(n_j, l_j):
        for m in range(2 * l + 1):
            if n == -1:
                n = '*'
            print('%2s %s %s_%s' % (i, n, 'spdf'[l], angular[l][m]))
            i += 1


def number_of_projectors(setup):
    """Returns the number of the bound state projectors.

    If setup is a string, treat this as an element name.
    """
    if isinstance(setup, str):
        setup = SetupData(setup, 'LDA', 'paw')
        n_j = setup.n_j
        l_j = setup.l_j
    else:
        n_j = setup.n_j
        l_j = setup.l_j

    i = 0
    for n, l in zip(n_j, l_j):
        for m in range(2 * l + 1):
            if n != -1:
                i += 1
    return i


def get_angular_projectors(setup, angular, type='bound'):
    """Determine the projector indices which have specified angula
    quantum number.

    angular can be s, p, d, f, or a list of these.
    If type is 'bound', only bound state projectors are considered, otherwise
    all projectors are included.
    """
    # Get the number of relevant j values
    if type == 'bound':
        nj = len([n for n in setup.n_j if n >= 0])
    else:
        nj = len(setup.n_j)

    # Choose the relevant projectors
    projectors = []
    i = 0
    j = 0
    for j in range(nj):
        m = 2 * setup.l_j[j] + 1
        if 'spdf'[setup.l_j[j]] in angular:
            projectors.extend(range(i, i + m))
        j += 1
        i += m

    return projectors


def delta(x, x0, width, mode='Gauss'):
    """Return a gaussian of given width centered at x0."""
    if mode == 'Gauss':
        return np.exp(np.clip(-((x - x0) / width)**2,
                              -100.0, 100.0)) / (sqrt(pi) * width)
    if mode == 'Lorentz':
        return (2 / pi / width) / ((np.clip(((x - x0) / (width / 2))**2,
                                            -100.0, 100.0)) + 1)


def fold(energies, weights, npts, width, mode='Gauss'):
    """Take a list of energies and weights, and sum a delta function
    for each."""
    emin = min(energies) - 5 * width
    emax = max(energies) + 5 * width
    e = np.linspace(emin, emax, npts)
    dos_e = np.zeros(npts)
    for e0, w in zip(energies, weights):
        dos_e += w * delta(e, e0, width, mode=mode)
    return e, dos_e


def raw_orbital_LDOS(paw, a, spin, angular='spdf', nbands=None):
    """Return a list of eigenvalues, and their weight on the specified atom.

    angular can be s, p, d, f, or a list of these.
    If angular is None, the raw weight for each projector is returned.

    An integer value for ``angular`` can also be used to specify a specific
    projector function.

    Setting nbands limits the number of bands included. This speeds up the
    calculation if one has many bands in the calculator but is only interested
    in the DOS at low energies.
    """
    wfs = paw.wfs
    w_k = wfs.kd.weight_k
    nk = len(w_k)
    if not nbands:
        nb = wfs.bd.nbands
    else:
        nb = nbands
        assert nb <= wfs.bd.nbands, ('nbands higher than available number' +
                                     'of bands')

    if a < 0:
        # Allow list-style negative indices; we'll need the positive a for the
        # dictionary lookup later
        a = len(wfs.setups) + a

    I1 = sum(setup.ni for setup in wfs.setups[:a])
    setup = wfs.setups[a]
    I2 = I1 + setup.ni

    energies = np.empty(nb * nk)
    weights_xi = np.empty((nb * nk, setup.ni))
    x = 0
    for k, w in enumerate(w_k):
        eps_n = wfs.collect_eigenvalues(k=k, s=spin)
        if len(eps_n) > 0:
            energies[x:x + nb] = eps_n[:nb]
        P_nI = wfs.collect_projections(k, spin)
        if P_nI is not None:
            weights_xi[x:x + nb, :] = w * abs(P_nI[:nb, I1:I2])**2
        x += nb

    wfs.world.broadcast(energies, 0)
    wfs.world.broadcast(weights_xi, 0)

    if angular is None:
        return energies, weights_xi
    elif isinstance(angular, int):
        return energies, weights_xi[:, angular]
    else:
        projectors = get_angular_projectors(setup, angular, type='bound')
        weights = np.sum(np.take(weights_xi,
                                 indices=projectors, axis=1), axis=1)
        return energies, weights


def all_electron_LDOS(paw, mol, spin, lc=None, wf_k=None, P_aui=None):
    """Returns a list of eigenvalues, and their weights on a given molecule

       wf_k should be a list of pseudo_wavefunctons of a Kohn-Sham state,
       corresponding to the different kpoints. It should be accompanied by a
       list of projector overlaps: P_aui=[[kpt.P_ani[a][n] for kpt in
       paw.wfs.kpt_u] for a in range(len(molecule))] for the band n. The
       weights are then calculated as the overlap of all-electron
       KS wavefunctions with wf_k

       If wf_k is None, the weights are calculated as linear combinations of
       atomic orbitals using P_uni. lc should then be a list of weights
       for each atom. For example, the pure 2pi_x orbital of a CO or N2
       molecule can be obtained with lc=[[0,0,0,1.0],[0,0,0,-1.0]]. mol
       should be a list of atom numbers contributing to the molecule."""

    w_k = paw.wfs.kd.weight_k
    nk = len(w_k)
    nb = paw.wfs.bd.nbands

    P_kn = np.zeros((nk, nb), complex)
    if wf_k is None:
        if lc is None:
            lc = [[1, 0, 0, 0] for a in mol]
        for k, kpt in enumerate(kpt_s[spin] for kpt_s in paw.wfs.kpt_qs):
            N = 0
            for atom, w_a in zip(mol, lc):
                i = 0
                for w_o in w_a:
                    P_kn[k] += w_o * kpt.P_ani[atom][:, i]
                    N += abs(w_o)**2
                    i += 1
        P_kn /= sqrt(N)

    else:
        P_aui = [np.array(P_ui).conj() for P_ui in P_aui]
        for k, kpt in enumerate(kpt_s[spin] for kpt_s in paw.wfs.kpt_qs):
            for n in range(nb):
                P_kn[k][n] = paw.wfs.integrate(wf_k[k], kpt.psit_nG[n])
                for a, b in zip(mol, range(len(mol))):
                    atom = paw.wfs.setups[a]
                    p_i = kpt.P_ani[a][n]
                    for i in range(len(p_i)):
                        for j in range(len(p_i)):
                            P_kn[k][n] += (P_aui[b][spin * nk + k][i] *
                                           atom.dO_ii[i][j] * p_i[j])
            print('# k', k, ' Sum_m |<m|n>|^2 =', sum(abs(P_kn[k])**2))

    energies = np.empty(nb * nk)
    weights = np.empty(nb * nk)
    x = 0
    for k, w in enumerate(w_k):
        eps_n = paw.wfs.collect_eigenvalues(k=k, s=spin)
        if len(eps_n) > 0:
            energies[x:x + nb] = eps_n
            weights[x:x + nb] = w * abs(P_kn[k])**2
        x += nb

    return energies, weights


def get_all_electron_IPR(paw):
    density = paw.density
    wfs = paw.wfs
    n_G = wfs.gd.empty()
    n_g = density.finegd.empty()
    print()
    print('inverse participation function')
    print('-' * 35)
    print('%5s %5s %10s %10s' % ('k', 'band', 'eps', 'ipr'))
    print('-' * 35)
    for k, kpt in enumerate(paw.wfs.kpt_u):
        for n, (eps, psit_G) in enumerate(zip(kpt.eps_n, kpt.psit_nG)):
            n_G[:] = 0.0
            wfs.add_orbital_density(n_G, kpt, n)
            density.interpolator.apply(n_G, n_g)
            norm = density.finegd.integrate(n_g)
            n_g = n_g ** 2
            ipr = density.finegd.integrate(n_g)
            for a in kpt.P_ani:
                # Get xccorr for atom a
                setup = paw.density.setups[a]
                xccorr = setup.xc_correction

                # Get D_sp for atom a
                D_sp = np.array(wfs.get_orbital_density_matrix(a, kpt, n))

                # density a function of L and partial wave radial pair
                # density coefficient
                D_sLq = gemmdot(D_sp, xccorr.B_Lqp, trans='t')

                # Create pseudo/ae density iterators for integration
                n_iter = xccorr.expand_density(D_sLq, xccorr.n_qg, None)
                nt_iter = xccorr.expand_density(D_sLq, xccorr.nt_qg, None)

                # Take the spherical average of smooth and ae radial
                # xc potentials
                for n_sg, nt_sg, integrator in zip(n_iter,
                                                   nt_iter,
                                                   xccorr.get_integrator(
                                                       None)):
                    ipr += integrator.weight * np.sum((n_sg[0]**2 -
                                                       nt_sg[0]**2) *
                                                      xccorr.rgd.dv_g)
                    norm += integrator.weight * np.sum((n_sg[0] - nt_sg[0]) *
                                                       xccorr.rgd.dv_g)

            print('%5i %5i %10.5f %10.5f' % (k, n, eps, ipr / norm**2))
    print('-' * 35)


def raw_wignerseitz_LDOS(paw, a, spin):
    """Return a list of eigenvalues, and their weight on the specified atom.
       Only supports real (gamma-point only) wavefunctions."""
    wfs = paw.wfs
    assert wfs.dtype == float
    gd = wfs.gd
    atom_index = wignerseitz(gd, paw.atoms)

    w_k = wfs.kd.weight_k
    nk = len(w_k)
    nb = wfs.bd.nbands

    energies = np.empty(nb * nk)
    weights = np.empty(nb * nk)
    x = 0
    for k, w in enumerate(w_k):
        u = spin * nk + k
        energies[x:x + nb] = wfs.collect_eigenvalues(k=k, s=spin)
        for n, psit_G in enumerate(wfs.kpt_u[u].psit_nG):
            P_i = wfs.kpt_u[u].P_ani[a][n]
            P_p = pack_density(np.outer(P_i, P_i))
            Delta_p = sqrt(4 * pi) * wfs.setups[a].Delta_pL[:, 0]
            weights[x + n] = w * (
                gd.integrate(abs(np.where(atom_index == a, psit_G, 0.0))**2) +
                np.dot(Delta_p, P_p))
        x += nb
    return energies, weights


class RawLDOS:
    """Class to get the unfolded LDOS"""
    def __init__(self, calc):
        self.paw = calc
        for setup in calc.wfs.setups.setups.values():
            if not hasattr(setup, 'l_i'):
                # get the mapping
                l_i = []
                for l in setup.l_j:
                    for m in range(2 * l + 1):
                        l_i.append(l)
                setup.l_i = l_i

    def get(self, atom):
        """Return the s,p,d weights for each state"""
        wfs = self.paw.wfs
        nibzkpts = len(wfs.kd.ibzk_kc)
        spd = np.zeros((wfs.nspins, nibzkpts, wfs.bd.nbands, 3))

        if hasattr(atom, '__iter__'):
            # atom is a list of atom indices
            for a in atom:
                spd += self.get(a)
            return spd

        l_i = wfs.setups[atom].l_i

        for kpt in self.paw.wfs.kpt_u:
            if atom in kpt.P_ani:
                for i, P_n in enumerate(kpt.P_ani[atom].T):
                    spd[kpt.s, kpt.k, :, l_i[i]] += np.abs(P_n)**2

        wfs.gd.comm.sum(spd)
        wfs.kd.comm.sum(spd)
        return spd

    def by_element(self, indices=None):
        """Return a dict with elements as keys and LDOS as values.

        Parameter:
        -----------
        indices: list or None
          Atom indices to consider, default None
        """
        atoms = self.paw.atoms
        if indices is None:
            selected = range(len(atoms))
        else:
            selected = indices

        elemi = {}
        for i in selected:
            symbol = atoms[i].symbol
            if symbol in elemi:
                elemi[symbol].append(i)
            else:
                elemi[symbol] = [i]
        for key in elemi:
            elemi[key] = self.get(elemi[key])
        return elemi

    def to_file(self,
                ldbe,
                filename,
                width=None,
                shift=True,
                bound=False):
        """Write the LDOS to a file.

        Parameters:
        -----------
        ldbe: dict
          weights
        filename: string
        width: float or None
          If a width is given, the LDOS will be Gaussian folded
        shift: bool
          Set Fermi energy to 0 eV. Default True
        bound: bool
          If the calculation was performed spin-polarized with fixmagmom=True,
          there are two Fermi-levels, one for each spin. False shifts
          both spins by their own Fermi levels, True both
          by the higher Fermi level. Default: False
        """

        f = paropen(filename, 'w')

        def append_weight_strings(ldbe, data):
            s = ''
            for key in ldbe:
                for l in 'spd':
                    data.append(key + '(' + l + ')-weight')
                if len(key) == 1:
                    key = ' ' + key
                s += ' ' + key + ':s     p        d       '
            return s

        wfs = self.paw.wfs

        if width is None:
            # unfolded ldos
            fmf = FMF(['Raw LDOS obtained from projector weights'])
            print(fmf.header(), end=' ', file=f)
            data = ['energy: energy [eV]',
                    'occupation number: occ',
                    'spin index: s',
                    'k-point index: k',
                    'band index: n',
                    'k-point weight: weight']
            string = '# e_i[eV]  occ     s     k      n   kptwght '
            string += append_weight_strings(ldbe, data)
            data.append('summed weights: sum')
            string += ' sum'
            print(fmf.data(data), end=' ', file=f)
            print(string, file=f)
            for k in range(wfs.kd.nibzkpts):
                for s in range(wfs.nspins):
                    e_n = self.paw.get_eigenvalues(kpt=k, spin=s)
                    f_n = self.paw.get_occupation_numbers(kpt=k, spin=s)
                    if e_n is None:
                        continue
                    w = wfs.kd.weight_k[k]
                    for n in range(wfs.bd.nbands):
                        sum = 0.0
                        print('%10.5f %6.4f %2d %5d' % (e_n[n], f_n[n],
                                                        s, k), end=' ', file=f)
                        print('%6d %8.4f' % (n, w), end=' ', file=f)
                        for key in ldbe:
                            spd = ldbe[key][s, k, n]
                            for l in range(3):
                                sum += spd[l]
                                print('%8.4f' % spd[l], end=' ', file=f)
                        print('%8.4f' % sum, file=f)
        else:
            # folded ldos
            fmf = FMF(['Folded raw LDOS obtained from projector weights'])
            print(fmf.header(), end=' ', file=f)

            gauss = Gauss(width)
            print(fmf.field('folding',
                            ['function: Gauss',
                             'width: ' + str(width) + ' [eV]']),
                  end=' ', file=f)

            data = ['energy: energy [eV]',
                    'spin index: s',
                    'k-point index: k',
                    'band index: n',
                    'k-point weight: weight']

            # minimal and maximal energies
            emin = 1.e32
            emax = -1.e32
            for k in range(wfs.kd.nibzkpts):
                for s in range(wfs.nspins):
                    e_n = self.paw.get_eigenvalues(kpt=k, spin=s)
                    emin = min(e_n.min(), emin)
                    emax = max(e_n.max(), emax)
            emin -= 4 * width
            emax += 4 * width

            # Fermi energies
            efermis = self.paw.wfs.fermi_levels

            eshift = 0.0

            if shift and len(efermis) == 1:
                eshift = -efermis[0]

            # set de to sample 4 points in the width
            de = width / 4

            string = '# e[eV]     s  '
            string += append_weight_strings(ldbe, data)

            for s in range(wfs.nspins):
                if len(efermis) == 2:
                    if not bound:
                        eshift = -efermis[s]
                    else:
                        eshift = -efermis.max()

                print(fmf.data(data), end=' ', file=f)

                print('# Gauss folded, width=%g [eV]' % width, file=f)
                if shift:
                    print('# shifted to Fermi energy = 0', file=f)
                    print('# Fermi energy was', end=' ', file=f)
                else:
                    print('# Fermi energy', end=' ', file=f)
                print(efermis, 'eV', file=f)
                print(string, file=f)

                # loop over energies
                emax = emax + 0.5 * de
                e = emin
                while e < emax:
                    val = {}
                    for key in ldbe:
                        val[key] = np.zeros(3)
                    for k in range(wfs.kd.nibzkpts):
                        w = wfs.kpt_u[k].weight
                        e_n = self.paw.get_eigenvalues(kpt=k, spin=s)
                        for n in range(wfs.bd.nbands):
                            w_i = w * gauss.get(e_n[n] - e)
                            for key in ldbe:
                                val[key] += w_i * ldbe[key][s, k, n]

                    print('%10.5f %2d' % (e + eshift, s), end=' ', file=f)
                    for key in val:
                        spd = val[key]
                        for l in range(3):
                            print('%8.4f' % spd[l], end=' ', file=f)
                    print(file=f)
                    e += de

        f.close()

    def by_element_to_file(self,
                           filename='ldos_by_element.dat',
                           width=None,
                           shift=True,
                           bound=False,
                           indices=None):
        """
        Write the LDOS by element to a file.

        Parameters:
        -----------
        indices: list or None
          Atom indices to consider, default None
        """
        ldbe = self.by_element(indices)
        self.to_file(ldbe, filename, width, shift, bound)


class LCAODOS:
    """Class for calculating the projected subspace DOS.

    The projected subspace density of states is defined only in LCAO
    mode.  The advantages to the PDOS based on projectors is that the
    LCAODOS will properly take non-orthogonality and completeness into
    account."""
    def __init__(self, calc):
        self.calc = calc

    def get_orbital_pdos(self, M, ravel=True):
        """Get projected DOS from LCAO basis function M."""
        return self.get_subspace_pdos([M], ravel=ravel)

    def get_atomic_subspace_pdos(self, a, ravel=True):
        """Get projected subspace DOS from LCAO basis on atom(s) a.

        a may be an atomic index or a list of indices."""
        Mvalues = self.get_atom_indices(a)
        return self.get_subspace_pdos(Mvalues, ravel=ravel)

    def get_atom_indices(self, a):
        """Get list of basis function indices of atom(s) a."""
        if isinstance(a, int):
            a = [a]
        Mvalues = []
        for a0 in a:
            M = self.calc.wfs.basis_functions.M_a[a0]
            Mvalues.extend(range(M, M + self.calc.wfs.setups[a0].nao))
        return Mvalues

    def get_subspace_pdos(self, Mvalues, ravel=True):
        """Get projected subspace DOS from LCAO basis."""
        wfs = self.calc.wfs
        bd = wfs.bd
        kd = wfs.kd
        gd = wfs.gd

        for kpt in wfs.kpt_u:
            assert not np.isnan(kpt.eps_n).any()

        w_skn = np.zeros((kd.nspins, kd.nibzkpts, bd.nbands))
        eps_skn = np.zeros((kd.nspins, kd.nibzkpts, bd.nbands))
        for u, kpt in enumerate(wfs.kpt_u):
            C_nM = kpt.C_nM
            from gpaw.kohnsham_layouts import BlacsOrbitalLayouts
            if isinstance(wfs.ksl, BlacsOrbitalLayouts):
                raise NotImplementedError('Something not quite working.  '
                                          'FIXME.')
                S_MM = wfs.ksl.mmdescriptor.collect_on_master(kpt.S_MM)
                if bd.rank != 0 or gd.rank != 0:
                    S_MM = np.empty((wfs.ksl.nao, wfs.ksl.nao),
                                    dtype=wfs.dtype)
                bd.comm.broadcast(S_MM, 0)
            else:
                S_MM = kpt.S_MM

            if gd.comm.rank == 0:
                S_mm = S_MM[Mvalues, :][:, Mvalues].copy()
                iS_mm = np.linalg.inv(S_mm).copy()
                CS_nm = np.dot(C_nM, S_MM[:, Mvalues])
                CSiS_nm = np.dot(CS_nm, iS_mm)
                w_n = (np.conj(CS_nm) * CSiS_nm).sum(axis=1) * kpt.weight
                w_skn[kpt.s, kpt.k, bd.beg:bd.end:bd.step] = w_n.real
                eps_skn[kpt.s, kpt.k, bd.beg:bd.end:bd.step] = kpt.eps_n

        for arr in [eps_skn, w_skn]:
            gd.comm.broadcast(arr, 0)
            bd.comm.sum(arr)
            kd.comm.sum(arr)

        if ravel:
            eps_n = eps_skn.ravel()
            w_n = w_skn.ravel()
            args = np.argsort(eps_n)
            eps_n = eps_n[args]
            w_n = w_n[args]
            return eps_n, w_n
        else:
            return eps_skn, w_skn


class RestartLCAODOS(LCAODOS):
    """Class for calculating LCAO subspace PDOS from a restarted calculator.

    Warning: This has side effects on the calculator.  The
    operation will allocate memory to diagonalize the Hamiltonian and
    set coefficients plus positions."""
    def __init__(self, calc):
        LCAODOS.__init__(self, calc)
        system = calc.get_atoms()
        calc.set_positions(system)
        calc.wfs.eigensolver.iterate(calc.hamiltonian, calc.wfs)
