import numpy as np
from ase.units import Ha

from gpaw.projections import Projections
from gpaw.utilities import pack, unpack2
from gpaw.utilities.blas import axpy, mmm
from gpaw.utilities.partition import AtomPartition


class WaveFunctions:
    """...

    setups:
        List of setup objects.
    symmetry:
        Symmetry object.
    kpt_u:
        List of **k**-point objects.
    nbands: int
        Number of bands.
    nspins: int
        Number of spins.
    dtype: dtype
        Data type of wave functions (float or complex).
    bzk_kc: ndarray
        Scaled **k**-points used for sampling the whole
        Brillouin zone - values scaled to [-0.5, 0.5).
    ibzk_kc: ndarray
        Scaled **k**-points in the irreducible part of the
        Brillouin zone.
    weight_k: ndarray
        Weights of the **k**-points in the irreducible part
        of the Brillouin zone (summing up to 1).
    kpt_comm:
        MPI-communicator for parallelization over **k**-points.
    """

    def __init__(self, gd, nvalence, setups, bd, dtype, collinear,
                 world, kd, kptband_comm, timer):
        self.gd = gd
        self.nspins = kd.nspins
        self.collinear = collinear
        self.nvalence = nvalence
        self.bd = bd
        self.dtype = dtype
        assert dtype == float or dtype == complex
        self.world = world
        self.kd = kd
        self.kptband_comm = kptband_comm
        self.timer = timer
        self.atom_partition = None

        self.kpt_qs = kd.create_k_points(self.gd.sdisp_cd, collinear)
        self.kpt_u = [kpt for kpt_s in self.kpt_qs for kpt in kpt_s]

        self.occupations = None
        self.fermi_levels = None

        self.eigensolver = None
        self.positions_set = False
        self.spos_ac = None

        self.set_setups(setups)

    @property
    def fermi_level(self):
        assert len(self.fermi_levels) == 1
        return self.fermi_levels[0]

    def summary(self, log):
        log(eigenvalue_string(self))

        if self.fermi_levels is None:
            return

        if len(self.fermi_levels) == 1:
            log(f'Fermi level: {self.fermi_levels[0] * Ha:.5f}\n')
        else:
            f1, f2 = (f * Ha for f in self.fermi_levels)
            log(f'Fermi levels: {f1:.5f}, {f2:.5f}\n')

    def set_setups(self, setups):
        self.setups = setups

    def set_eigensolver(self, eigensolver):
        self.eigensolver = eigensolver

    def add_realspace_orbital_to_density(self, nt_G, psit_G):
        if psit_G.dtype == float:
            axpy(1.0, psit_G**2, nt_G)
        else:
            assert psit_G.dtype == complex
            axpy(1.0, psit_G.real**2, nt_G)
            axpy(1.0, psit_G.imag**2, nt_G)

    def add_orbital_density(self, nt_G, kpt, n):
        self.add_realspace_orbital_to_density(nt_G, kpt.psit_nG[n])

    def calculate_band_energy(self):
        e_band = 0.0
        for kpt in self.kpt_u:
            e_band += np.dot(kpt.f_n, kpt.eps_n)

        try:  # DCSF needs this ...
            e_band += self.occupations.calculate_band_energy(self)
        except AttributeError:
            pass

        return self.kptband_comm.sum(e_band)

    def calculate_density_contribution(self, nt_sG):
        """Calculate contribution to pseudo density from wave functions.

        Array entries are written to (not added to)."""
        nt_sG.fill(0.0)
        for kpt in self.kpt_u:
            self.add_to_density_from_k_point(nt_sG, kpt)
        self.kptband_comm.sum(nt_sG)

        self.timer.start('Symmetrize density')
        for nt_G in nt_sG:
            self.kd.symmetry.symmetrize(nt_G, self.gd)
        self.timer.stop('Symmetrize density')

    def add_to_density_from_k_point(self, nt_sG, kpt):
        self.add_to_density_from_k_point_with_occupation(nt_sG, kpt, kpt.f_n)

    def get_orbital_density_matrix(self, a, kpt, n):
        """Add the nth band density from kpt to density matrix D_sp"""
        ni = self.setups[a].ni
        D_sii = np.zeros((self.nspins, ni, ni))
        P_i = kpt.P_ani[a][n]
        D_sii[kpt.s] += np.outer(P_i.conj(), P_i).real
        D_sp = [pack(D_ii) for D_ii in D_sii]
        return D_sp

    def calculate_atomic_density_matrices_k_point(self, D_sii, kpt, a, f_n):
        if kpt.rho_MM is not None:
            P_Mi = self.P_aqMi[a][kpt.q]
            rhoP_Mi = np.zeros_like(P_Mi)
            D_ii = np.zeros(D_sii[kpt.s].shape, kpt.rho_MM.dtype)
            mmm(1.0, kpt.rho_MM, 'N', P_Mi, 'N', 0.0, rhoP_Mi)
            mmm(1.0, P_Mi, 'C', rhoP_Mi, 'N', 0.0, D_ii)
            D_sii[kpt.s] += D_ii.real
        else:
            if self.collinear:
                P_ni = kpt.projections[a]
                D_sii[kpt.s] += np.dot(P_ni.T.conj() * f_n, P_ni).real
            else:
                P_nsi = kpt.projections[a]
                D_ssii = np.einsum('nsi,n,nzj->szij',
                                   P_nsi.conj(), f_n, P_nsi)
                D_sii[0] += (D_ssii[0, 0] + D_ssii[1, 1]).real
                D_sii[1] += 2 * D_ssii[0, 1].real
                D_sii[2] += 2 * D_ssii[0, 1].imag
                D_sii[3] += (D_ssii[0, 0] - D_ssii[1, 1]).real

        if hasattr(kpt, 'c_on'):
            for ne, c_n in zip(kpt.ne_o, kpt.c_on):
                d_nn = ne * np.outer(c_n.conj(), c_n)
                D_sii[kpt.s] += np.dot(P_ni.T.conj(), np.dot(d_nn, P_ni)).real

    def calculate_atomic_density_matrices(self, D_asp):
        """Calculate atomic density matrices from projections."""
        f_un = [kpt.f_n for kpt in self.kpt_u]
        self.calculate_atomic_density_matrices_with_occupation(D_asp, f_un)

    def calculate_atomic_density_matrices_with_occupation(self, D_asp, f_un):
        """Calculate atomic density matrices from projections with
        custom occupation f_un."""

        # Parameter check (if user accidentally passes f_n instead of f_un)
        if f_un[0] is not None:  # special case for transport calculations...
            assert isinstance(f_un[0], np.ndarray)
        # Varying f_n used in calculation of response part of GLLB-potential
        for a, D_sp in D_asp.items():
            ni = self.setups[a].ni
            D_sii = np.zeros((len(D_sp), ni, ni))
            for f_n, kpt in zip(f_un, self.kpt_u):
                self.calculate_atomic_density_matrices_k_point(D_sii, kpt, a,
                                                               f_n)
            D_sp[:] = [pack(D_ii) for D_ii in D_sii]
            self.kptband_comm.sum(D_sp)

        self.symmetrize_atomic_density_matrices(D_asp)

    def symmetrize_atomic_density_matrices(self, D_asp):
        if len(self.kd.symmetry.op_scc) == 0:
            return

        a_sa = self.kd.symmetry.a_sa
        D_asp.redistribute(self.atom_partition.as_serial())
        for s in range(self.nspins):
            D_aii = [unpack2(D_asp[a][s])
                     for a in range(len(D_asp))]
            for a, D_ii in enumerate(D_aii):
                setup = self.setups[a]
                D_asp[a][s] = pack(setup.symmetrize(a, D_aii, a_sa))
        D_asp.redistribute(self.atom_partition)

    def calculate_occupation_numbers(self, fixed_fermi_level=False):
        if self.collinear and self.nspins == 1:
            degeneracy = 2
        else:
            degeneracy = 1

        f_qn, fermi_levels, e_entropy = self.occupations.calculate(
            nelectrons=self.nvalence / degeneracy,
            eigenvalues=[kpt.eps_n * Ha for kpt in self.kpt_u],
            weights=[kpt.weightk for kpt in self.kpt_u],
            fermi_levels_guess=self.fermi_levels * Ha
            if self.fermi_levels is not None else None)

        if not fixed_fermi_level or self.fermi_levels is None:
            self.fermi_levels = np.array(fermi_levels) / Ha

        for f_n, kpt in zip(f_qn, self.kpt_u):
            kpt.f_n = f_n * (kpt.weightk * degeneracy)

        return e_entropy * degeneracy / Ha

    def set_positions(self, spos_ac, atom_partition=None):
        self.positions_set = False
        # rank_a = self.gd.get_ranks_from_positions(spos_ac)
        # atom_partition = AtomPartition(self.gd.comm, rank_a)
        # XXX pass AtomPartition around instead of spos_ac?
        # All the classes passing around spos_ac end up needing the ranks
        # anyway.

        if atom_partition is None:
            rank_a = self.gd.get_ranks_from_positions(spos_ac)
            atom_partition = AtomPartition(self.gd.comm, rank_a)

        if self.atom_partition is not None and self.kpt_u[0].P_ani is not None:
            with self.timer('Redistribute'):
                for kpt in self.kpt_u:
                    P = kpt.projections
                    assert self.atom_partition == P.atom_partition
                    kpt.projections = P.redist(atom_partition)
                    assert atom_partition == kpt.projections.atom_partition

        self.atom_partition = atom_partition
        self.kd.symmetry.check(spos_ac)
        self.spos_ac = spos_ac

    def allocate_arrays_for_projections(self, my_atom_indices):  # XXX unused
        if not self.positions_set and self.kpt_u[0]._projections is not None:
            # Projections have been read from file - don't delete them!
            pass
        else:
            nproj_a = [setup.ni for setup in self.setups]
            for kpt in self.kpt_u:
                kpt.projections = Projections(
                    self.bd.nbands, nproj_a,
                    self.atom_partition,
                    self.bd.comm,
                    collinear=self.collinear, spin=kpt.s, dtype=self.dtype)

    def collect_eigenvalues(self, k, s):
        return self.collect_array('eps_n', k, s)

    def collect_occupations(self, k, s):
        return self.collect_array('f_n', k, s)

    def collect_array(self, name, k, s, subset=None):
        """Helper method for collect_eigenvalues and collect_occupations.

        For the parallel case find the rank in kpt_comm that contains
        the (k,s) pair, for this rank, collect on the corresponding
        domain a full array on the domain master and send this to the
        global master."""

        kpt_qs = self.kpt_qs
        kpt_rank, q = self.kd.get_rank_and_index(k)
        if self.kd.comm.rank == kpt_rank:
            a_nx = getattr(kpt_qs[q][s], name)

            if subset is not None:
                a_nx = a_nx[subset]

            # Domain master send this to the global master
            if self.gd.comm.rank == 0:
                if self.bd.comm.size == 1:
                    if kpt_rank == 0:
                        return a_nx
                    else:
                        self.kd.comm.ssend(a_nx, 0, 1301)
                else:
                    b_nx = self.bd.collect(a_nx)
                    if self.bd.comm.rank == 0:
                        if kpt_rank == 0:
                            return b_nx
                        else:
                            self.kd.comm.ssend(b_nx, 0, 1301)

        elif self.world.rank == 0 and kpt_rank != 0:
            # Only used to determine shape and dtype of receiving buffer:
            a_nx = getattr(kpt_qs[0][0], name)

            if subset is not None:
                a_nx = a_nx[subset]

            b_nx = np.zeros((self.bd.nbands,) + a_nx.shape[1:],
                            dtype=a_nx.dtype)
            self.kd.comm.receive(b_nx, kpt_rank, 1301)
            return b_nx

        return np.zeros(0)  # see comment in get_wave_function_array() method

    def collect_auxiliary(self, value, k, s, shape=1, dtype=float):
        """Helper method for collecting band-independent scalars/arrays.

        For the parallel case find the rank in kpt_comm that contains
        the (k,s) pair, for this rank, collect on the corresponding
        domain a full array on the domain master and send this to the
        global master."""

        kpt_rank, q = self.kd.get_rank_and_index(k)

        if self.kd.comm.rank == kpt_rank:
            if isinstance(value, str):
                a_o = getattr(self.kpt_qs[q][s], value)
            else:
                u = q * self.nspins + s
                a_o = value[u]  # assumed list

            # Make sure data is a mutable object
            a_o = np.asarray(a_o)

            if a_o.dtype is not dtype:
                a_o = a_o.astype(dtype)

            # Domain master send this to the global master
            if self.gd.comm.rank == 0:
                if kpt_rank == 0:
                    return a_o
                else:
                    self.kd.comm.send(a_o, 0, 1302)

        elif self.world.rank == 0 and kpt_rank != 0:
            b_o = np.zeros(shape, dtype=dtype)
            self.kd.comm.receive(b_o, kpt_rank, 1302)
            return b_o

    def collect_projections(self, k, s):
        """Helper method for collecting projector overlaps across domains.

        For the parallel case find the rank in kpt_comm that contains
        the (k,s) pair, for this rank, send to the global master."""

        kpt_rank, q = self.kd.get_rank_and_index(k)

        if self.kd.comm.rank == kpt_rank:
            kpt = self.kpt_qs[q][s]
            P_nI = kpt.projections.collect()
            if self.world.rank == 0:
                return P_nI
            if P_nI is not None:
                self.kd.comm.send(np.ascontiguousarray(P_nI), 0, tag=117)
        if self.world.rank == 0:
            nproj = sum(setup.ni for setup in self.setups)
            if not self.collinear:
                nproj *= 2
            P_nI = np.empty((self.bd.nbands, nproj), self.dtype)
            self.kd.comm.receive(P_nI, kpt_rank, tag=117)
            return P_nI

    def get_wave_function_array(self, n, k, s, realspace=True, periodic=False):
        """Return pseudo-wave-function array on master.

        n: int
            Global band index.
        k: int
            Global IBZ k-point index.
        s: int
            Spin index (0 or 1).
        realspace: bool
            Transform plane wave or LCAO expansion coefficients to real-space.

        For the parallel case find the ranks in kd.comm and bd.comm
        that contains to (n, k, s), and collect on the corresponding
        domain a full array on the domain master and send this to the
        global master."""

        kpt_rank, q = self.kd.get_rank_and_index(k)
        band_rank, myn = self.bd.who_has(n)

        rank = self.world.rank

        if (self.kd.comm.rank == kpt_rank and
            self.bd.comm.rank == band_rank):
            u = q * self.nspins + s
            psit_G = self._get_wave_function_array(u, myn,
                                                   realspace, periodic)

            if realspace:
                psit_G = self.gd.collect(psit_G)

            if rank == 0:
                return psit_G

            # Domain master send this to the global master
            if self.gd.comm.rank == 0:
                psit_G = np.ascontiguousarray(psit_G)
                self.world.ssend(psit_G, 0, 1398)

        if rank == 0:
            # allocate full wave function and receive
            shape = () if self.collinear else (2,)
            psit_G = self.empty(shape, global_array=True,
                                realspace=realspace)
            # XXX this will fail when using non-standard nesting
            # of communicators.
            world_rank = (kpt_rank * self.gd.comm.size *
                          self.bd.comm.size +
                          band_rank * self.gd.comm.size)
            self.world.receive(psit_G, world_rank, 1398)
            return psit_G

        # We return a number instead of None on all the slaves.  Most of
        # the time the return value will be ignored on the slaves, but
        # in some cases it will be multiplied by some other number and
        # then ignored.  Allowing for this will simplify some code here
        # and there.
        return np.nan

    def get_homo_lumo(self, spin=None):
        """Return HOMO and LUMO eigenvalues."""
        if spin is None:
            if self.nspins == 1:
                return self.get_homo_lumo(0)
            h0, l0 = self.get_homo_lumo(0)
            h1, l1 = self.get_homo_lumo(1)
            return np.array([max(h0, h1), min(l0, l1)])

        n = self.nvalence // 2
        band_rank, myn = self.bd.who_has(n - 1)
        homo = -np.inf
        if self.bd.comm.rank == band_rank:
            for kpt in self.kpt_u:
                if kpt.s == spin:
                    homo = max(kpt.eps_n[myn], homo)
        homo = self.world.max(homo)

        lumo = np.inf
        if n < self.bd.nbands:  # there are not enough bands for LUMO
            band_rank, myn = self.bd.who_has(n)
            if self.bd.comm.rank == band_rank:
                for kpt in self.kpt_u:
                    if kpt.s == spin:
                        lumo = min(kpt.eps_n[myn], lumo)
            lumo = self.world.min(lumo)

        return np.array([homo, lumo])

    def write(self, writer):
        writer.write(version=2, ha=Ha)
        writer.write(fermi_levels=self.fermi_levels * Ha)
        writer.write(kpts=self.kd)
        self.write_projections(writer)
        self.write_eigenvalues(writer)
        self.write_occupations(writer)

    def write_projections(self, writer):
        nproj = sum(setup.ni for setup in self.setups)

        if self.collinear:
            shape = (self.nspins, self.kd.nibzkpts, self.bd.nbands, nproj)
        else:
            shape = (self.kd.nibzkpts, self.bd.nbands, 2, nproj)

        writer.add_array('projections', shape, self.dtype)

        for s in range(self.nspins):
            for k in range(self.kd.nibzkpts):
                P_nI = self.collect_projections(k, s)
                if not self.collinear and P_nI is not None:
                    P_nI.shape = (self.bd.nbands, 2, nproj)
                writer.fill(P_nI)

    def write_eigenvalues(self, writer):
        if self.collinear:
            shape = (self.nspins, self.kd.nibzkpts, self.bd.nbands)
        else:
            shape = (self.kd.nibzkpts, self.bd.nbands)

        writer.add_array('eigenvalues', shape)
        for s in range(self.nspins):
            for k in range(self.kd.nibzkpts):
                writer.fill(self.collect_eigenvalues(k, s) * Ha)

    def write_occupations(self, writer):

        if self.collinear:
            shape = (self.nspins, self.kd.nibzkpts, self.bd.nbands)
        else:
            shape = (self.kd.nibzkpts, self.bd.nbands)

        writer.add_array('occupations', shape)
        for s in range(self.nspins):
            for k in range(self.kd.nibzkpts):
                # Scale occupation numbers when writing:
                # XXX fix this in the code also ...
                weight = self.kd.weight_k[k] * 2 / self.nspins
                writer.fill(self.collect_occupations(k, s) / weight)

    def read(self, reader):
        r = reader.wave_functions
        # Backward compatibility:
        # Take parameters from main reader
        if 'ha' not in r:
            r.ha = reader.ha
        if 'version' not in r:
            r.version = reader.version

        if reader.version >= 3:
            self.fermi_levels = r.fermi_levels / r.ha
        else:
            o = reader.occupations
            self.fermi_levels = np.array(
                [o.fermilevel + o.split / 2,
                 o.fermilevel - o.split / 2]) / r.ha
            if self.occupations.name != 'fixmagmom':
                assert o.split == 0.0
                self.fermi_levels = self.fermi_levels[:1]

        if reader.version >= 2:
            kpts = r.kpts
            assert np.allclose(kpts.ibzkpts, self.kd.ibzk_kc)
            assert np.allclose(kpts.bzkpts, self.kd.bzk_kc)
            assert (kpts.bz2ibz == self.kd.bz2ibz_k).all()
            assert np.allclose(kpts.weights, self.kd.weight_k)

        self.read_projections(r)
        self.read_eigenvalues(r, r.version <= 0)
        self.read_occupations(r, r.version <= 0)

    def read_projections(self, reader):
        nslice = self.bd.get_slice()
        nproj_a = [setup.ni for setup in self.setups]
        atom_partition = AtomPartition(self.gd.comm,
                                       np.zeros(len(nproj_a), int))
        for u, kpt in enumerate(self.kpt_u):
            if self.collinear:
                index = (kpt.s, kpt.k)
            else:
                index = (kpt.k,)
            kpt.projections = Projections(
                self.bd.nbands, nproj_a,
                atom_partition, self.bd.comm,
                collinear=self.collinear, spin=kpt.s, dtype=self.dtype)
            if self.gd.comm.rank == 0:
                P_nI = reader.proxy('projections', *index)[nslice]
                if not self.collinear:
                    P_nI.shape = (self.bd.mynbands, -1)
                kpt.projections.matrix.array[:] = P_nI

    def read_eigenvalues(self, reader, old=False):
        nslice = self.bd.get_slice()
        for u, kpt in enumerate(self.kpt_u):
            if self.collinear:
                index = (kpt.s, kpt.k)
            else:
                index = (kpt.k,)
            eps_n = reader.proxy('eigenvalues', *index)[nslice]
            x = self.bd.mynbands - len(eps_n)  # missing bands?
            if x > 0:
                # Working on a real fix to this parallelization problem ...
                eps_n = np.pad(eps_n, (0, x), 'constant')
            if not old:  # skip for old tar-files gpw's
                eps_n /= reader.ha
            kpt.eps_n = eps_n

    def read_occupations(self, reader, old=False):
        nslice = self.bd.get_slice()
        for u, kpt in enumerate(self.kpt_u):
            if self.collinear:
                index = (kpt.s, kpt.k)
            else:
                index = (kpt.k,)
            f_n = reader.proxy('occupations', *index)[nslice]
            x = self.bd.mynbands - len(f_n)  # missing bands?
            if x > 0:
                # Working on a real fix to this parallelization problem ...
                f_n = np.pad(f_n, (0, x), 'constant')
            if not old:  # skip for old tar-files gpw's
                f_n *= kpt.weight
            kpt.f_n = f_n


def eigenvalue_string(wfs, comment=' '):
    """Write eigenvalues and occupation numbers into a string.

    The parameter comment can be used to comment out non-numers,
    for example to escape it for gnuplot.
    """
    tokens = []

    def add(*line):
        for token in line:
            tokens.append(token)
        tokens.append('\n')

    def eigs(k, s):
        eps_n = wfs.collect_eigenvalues(k, s)
        return eps_n * Ha

    def occs(k, s):
        occ_n = wfs.collect_occupations(k, s)
        return occ_n / wfs.kd.weight_k[k]

    if len(wfs.kd.ibzk_kc) == 1:
        if wfs.nspins == 1:
            add(comment, 'Band  Eigenvalues  Occupancy')
            eps_n = eigs(0, 0)
            f_n = occs(0, 0)
            if wfs.world.rank == 0:
                for n in range(wfs.bd.nbands):
                    add('%5d  %11.5f  %9.5f' % (n, eps_n[n], f_n[n]))
        else:
            add(comment, '                  Up                     Down')
            add(comment, 'Band  Eigenvalues  Occupancy  Eigenvalues  '
                'Occupancy')
            epsa_n = eigs(0, 0)
            epsb_n = eigs(0, 1)
            fa_n = occs(0, 0)
            fb_n = occs(0, 1)
            if wfs.world.rank == 0:
                for n in range(wfs.bd.nbands):
                    add('%5d  %11.5f  %9.5f  %11.5f  %9.5f' %
                        (n, epsa_n[n], fa_n[n], epsb_n[n], fb_n[n]))
        return ''.join(tokens)

    if len(wfs.kd.ibzk_kc) > 2:
        add('Showing only first 2 kpts')
        print_range = 2
    else:
        add('Showing all kpts')
        print_range = len(wfs.kd.ibzk_kc)

    if wfs.nvalence / 2. > 2:
        m = int(wfs.nvalence / 2. - 2)
    else:
        m = 0
    if wfs.bd.nbands - wfs.nvalence / 2. > 2:
        j = int(wfs.nvalence / 2. + 2)
    else:
        j = int(wfs.bd.nbands)

    if wfs.nspins == 1:
        add(comment, 'Kpt  Band  Eigenvalues  Occupancy')
        for i in range(print_range):
            eps_n = eigs(i, 0)
            f_n = occs(i, 0)
            if wfs.world.rank == 0:
                for n in range(m, j):
                    add('%3i %5d  %11.5f  %9.5f' % (i, n, eps_n[n], f_n[n]))
                add()
    else:
        add(comment, '                     Up                     Down')
        add(comment, 'Kpt  Band  Eigenvalues  Occupancy  Eigenvalues  '
            'Occupancy')

        for i in range(print_range):
            epsa_n = eigs(i, 0)
            epsb_n = eigs(i, 1)
            fa_n = occs(i, 0)
            fb_n = occs(i, 1)
            if wfs.world.rank == 0:
                for n in range(m, j):
                    add('%3i %5d  %11.5f  %9.5f  %11.5f  %9.5f' %
                        (i, n, epsa_n[n], fa_n[n], epsb_n[n], fb_n[n]))
                add()
    return ''.join(tokens)
