"""Kohn-Sham single particle excitations realated objects.

"""
import sys
import json
import numpy as np

from ase.units import Bohr, Hartree, alpha

import gpaw.mpi as mpi
from gpaw.utilities import packed_index
from gpaw.lrtddft.excitation import Excitation, ExcitationList, get_filehandle
from gpaw.pair_density import PairDensity
from gpaw.fd_operators import Gradient
from gpaw.utilities.tools import coordinates


class KSSingles(ExcitationList):
    """Kohn-Sham single particle excitations

    Input parameters:

    calculator:
      the calculator object after a ground state calculation

    nspins:
      number of spins considered in the calculation
      Note: Valid only for unpolarised ground state calculation

    eps:
      Minimal occupation difference for a transition (default 0.001)

    istart:
      First occupied state to consider
    jend:
      Last unoccupied state to consider
    energy_range:
      The energy range [emin, emax] or emax for KS transitions to use as basis
    """
    def __init__(self,
                 restrict={},
                 log=None,
                 txt=None):
        ExcitationList.__init__(self, log=log, txt=txt)
        self.world = mpi.world

        self.restrict = KSSRestrictor()
        self.restrict.update(restrict)

    def calculate(self, atoms, nspins=None):
        calculator = atoms.calc
        self.calculator = calculator

        # LCAO calculation requires special actions
        self.lcao = calculator.parameters.mode == 'lcao'

        # deny hybrids as their empty states are wrong
#        gsxc = calculator.hamiltonian.xc
#        hybrid = hasattr(gsxc, 'hybrid') and gsxc.hybrid > 0.0
#        assert(not hybrid)

        # ensure correctly initialized wave functions
        calculator.converge_wave_functions()
        self.world = calculator.wfs.world

        # parallelization over bands not yet supported
        assert(calculator.wfs.bd.comm.size == 1)

        # do the evaluation
        self.select(nspins)

        trkm = self.get_trk()
        self.log('KSS {0} transitions (restrict={1})'.format(
            len(self), self.restrict))
        self.log('KSS TRK sum %g (%g,%g,%g)' %
                 (np.sum(trkm) / 3., trkm[0], trkm[1], trkm[2]))
        pol = self.get_polarizabilities(lmax=3)
        self.log('KSS polarisabilities(l=0-3) %g, %g, %g, %g' %
                 tuple(pol.tolist()))
        return self

    @staticmethod
    def emin_emax(energy_range):
        emin = -sys.float_info.max
        emax = sys.float_info.max
        if energy_range is not None:
            try:
                emin, emax = energy_range
                emin /= Hartree
                emax /= Hartree
            except TypeError:
                emax = energy_range / Hartree
        return emin, emax

    def select(self, nspins=None):
        """Select KSSingles according to the given criterium."""

        # criteria
        emin, emax = self.restrict.emin_emax()
        istart = self.restrict['istart']
        jend = self.restrict['jend']
        eps = self.restrict['eps']

        if not hasattr(self, 'calculator'):  # I'm read from a file
            # throw away all not needed entries
            for i, ks in reversed(list(enumerate(self))):
                if ((ks.fij / ks.weight) <= eps or
                    ks.i < istart or ks.j > jend or
                    ks.energy < emin or ks.energy > emax):
                    del(self[i])
            return None

        paw = self.calculator
        wfs = paw.wfs
        self.dtype = wfs.dtype
        self.kpt_u = wfs.kpt_u

        if not self.lcao and self.kpt_u[0].psit_nG is None:
            raise RuntimeError('No wave functions in calculator!')

        # here, we need to take care of the spins also for
        # closed shell systems (Sz=0)
        # vspin is the virtual spin of the wave functions,
        #       i.e. the spin used in the ground state calculation
        # pspin is the physical spin of the wave functions
        #       i.e. the spin of the excited states
        self.nvspins = wfs.nspins
        self.npspins = wfs.nspins
        fijscale = 1
        ispins = [0]
        nks = wfs.kd.nibzkpts * wfs.kd.nspins
        if self.nvspins < 2:
            if (nspins or 0) > self.nvspins:
                self.npspins = nspins
                fijscale = 0.5
                ispins = [0, 1]
                nks *= 2

        kpt_comm = self.calculator.wfs.kd.comm
        nbands = len(self.kpt_u[0].f_n)

        # select
        take = np.zeros((nks, nbands, nbands), dtype=int)
        u = 0
        for ispin in ispins:
            for k in range(wfs.kd.nibzkpts):
                q = k - wfs.kd.k0
                for s in range(wfs.nspins):
                    if q >= 0 and q < wfs.kd.mynk:
                        kpt = wfs.kpt_qs[q][s]
                        for i in range(nbands):
                            for j in range(i + 1, nbands):
                                fij = (kpt.f_n[i] - kpt.f_n[j]) / kpt.weight
                                epsij = kpt.eps_n[j] - kpt.eps_n[i]
                                if (fij > eps and
                                    epsij >= emin and epsij < emax and
                                        i >= istart and j <= jend):
                                    take[u, i, j] = 1
                    u += 1
        kpt_comm.sum(take)

        self.log()
        self.log('Kohn-Sham single transitions')
        self.log()

        # calculate in parallel
        u = 0
        for ispin in ispins:
            for k in range(wfs.kd.nibzkpts):
                q = k - wfs.kd.k0
                for s in range(wfs.kd.nspins):
                    for i in range(nbands):
                        for j in range(i + 1, nbands):
                            if take[u, i, j]:
                                if q >= 0 and q < wfs.kd.mynk:
                                    kpt = wfs.kpt_qs[q][s]
                                    pspin = max(kpt.s, ispin)
                                    self.append(
                                        KSSingle(i, j, pspin, kpt, paw,
                                                 fijscale=fijscale,
                                                 dtype=self.dtype))
                                else:
                                    self.append(KSSingle(i, j, pspin=0,
                                                         kpt=None, paw=paw,
                                                         dtype=self.dtype))
                    u += 1

        # distribute
        for kss in self:
            kss.distribute()

    @classmethod
    def read(cls, filename=None, fh=None, restrict={}, log=None):
        """Read myself from a file"""
        assert (filename is not None) or (fh is not None)
        
        def fail(f):
            raise RuntimeError(f.name + ' does not contain ' +
                               cls.__class__.__name__ + ' data')
        if fh is None:
            f = get_filehandle(cls, filename)

            # there can be other information, i.e. the LrTDDFT header
            try:
                content = f.read()
                f.seek(content.index('# KSSingles'))
                del(content)
                f.readline()
            except ValueError:
                fail(f)
        else:
            f = fh
            # we assume to be at the right place and read the header
            if not f.readline().strip() == '# KSSingles':
                fail(f)

        words = f.readline().split()
        n = int(words[0])
        kssl = cls(log=log)
        if len(words) == 1:
            # very old output style for real wave functions (finite systems)
            kssl.dtype = float
            restrict_from_file = {}
        else:
            if words[1].startswith('complex'):
                kssl.dtype = complex
            else:
                kssl.dtype = float
            restrict_from_file = json.loads(f.readline())
            if not isinstance(restrict_from_file, dict):  # old output style
                restrict_from_file = {'eps': restrict_from_file}
        kssl.npspins = 1
        for i in range(n):
            kss = KSSingle(string=f.readline(), dtype=kssl.dtype)
            kssl.append(kss)
            kssl.npspins = max(kssl.npspins, kss.pspin + 1)
        
        if fh is None:
            f.close()

        kssl.update()
        kssl.restrict.update(restrict_from_file)
        if len(restrict):
            kssl.restrict.update(restrict)
            kssl.select()
        
        return kssl

    def update(self):
        istart = self[0].i
        jend = 0
        npspins = 1
        nvspins = 1
        for kss in self:
            istart = min(kss.i, istart)
            jend = max(kss.j, jend)
            if kss.pspin == 1:
                npspins = 2
            if kss.spin == 1:
                nvspins = 2
        self.restrict.update({'istart': istart, 'jend': jend})
        self.npspins = npspins
        self.nvspins = nvspins

        if hasattr(self, 'energies'):
            del(self.energies)

    def set_arrays(self):
        if hasattr(self, 'energies'):
            return
        energies = []
        fij = []
        me = []
        mur = []
        muv = []
        magn = []
        for k in self:
            energies.append(k.energy)
            fij.append(k.fij)
            me.append(k.me)
            mur.append(k.mur)
            if k.muv is not None:
                muv.append(k.muv)
            if k.magn is not None:
                magn.append(k.magn)
        self.energies = np.array(energies)
        self.fij = np.array(fij)
        self.me = np.array(me)
        self.mur = np.array(mur)
        if len(muv):
            self.muv = np.array(muv)
        else:
            self.muv = None
        if len(magn):
            self.magn = np.array(magn)
        else:
            self.magn = None

    def write(self, filename=None, fh=None):
        """Write current state to a file.

        'filename' is the filename. If the filename ends in .gz,
        the file is automatically saved in compressed gzip format.

        'fh' is a filehandle. This can be used to write into already
        opened files.
        """
        if self.world.rank != 0:
            return

        if fh is None:
            f = get_filehandle(self, filename, mode='w')
        else:
            f = fh

        f.write('# KSSingles\n')
        f.write('{0} {1}\n'.format(len(self), np.dtype(self.dtype)))
        f.write(json.dumps(self.restrict.values) + '\n')
        for kss in self:
            f.write(kss.outstring())
        if fh is None:
            f.close()

    def overlap(self, ov_nn, other):
        """Matrix element overlaps determined from wave function overlaps.

        Parameters
        ----------
        ov_nn: array
            Wave function overlap factors from a displaced calculator.
            Index 0 corresponds to our own wavefunctions conjugated and
            index 1 to the others' wavefunctions

        Returns
        -------
        ov_pp: array
            Overlap corresponding to matrix elements.
            Index 0 corresponds to our own matrix elements conjugated and
            index 1 to the others' matrix elements
        """
        n0 = len(self)
        n1 = len(other)
        ov_pp = np.zeros((n0, n1), dtype=ov_nn.dtype)
        i1_p = [ex.i for ex in other]
        j1_p = [ex.j for ex in other]
        for p0, ex0 in enumerate(self):
            ov_pp[p0, :] = ov_nn[ex0.i, i1_p].conj() * ov_nn[ex0.j, j1_p]
        return ov_pp


class KSSRestrictor:
    """Object to handle KSSingles restrictions"""
    defaults = {'eps': 0.01,
                'istart': 0,
                'jend': sys.maxsize,
                'energy_range': None}
    
    def __init__(self, dictionary={}):
        self._vals = {}
        self.update(dictionary)

    def __getitem__(self, index):
        assert index in self.defaults
        return self._vals.get(index, self.defaults[index])
    
    def __setitem__(self, index, value):
        assert index in self.defaults
        self._vals[index] = value
        
    def update(self, dictionary):
        for key, value in dictionary.items():
            self[key] = value
            
    def emin_emax(self):
        emin = -sys.float_info.max
        emax = sys.float_info.max
        if self['energy_range'] is not None:
            try:
                emin, emax = self['energy_range']
                emin /= Hartree
                emax /= Hartree
            except TypeError:
                emax = self['energy_range'] / Hartree
        return emin, emax

    def __ge__(self, other):
        """am I less or equal restricting than the other?"""
        emin, emax = self.emin_emax()
        omin, omax = other.emin_emax()
        res = (emin <= omin) & (emax >= omax)
        
        res &= self['istart'] <= other['istart']
        res &= self['jend'] >= other['jend']
        res &= self['eps'] < other['eps']

        return res

    @property
    def values(self):
        return dict(self._vals)

    def __str__(self):
        return str(self.values)


class KSSingle(Excitation, PairDensity):
    """Single Kohn-Sham transition containing all it's indicees

      pspin=physical spin
      spin=virtual  spin, i.e. spin in the ground state calc.
      kpt=the Kpoint object
      fijscale=weight for the occupation difference::
      me  = sqrt(fij*epsij) * <i|r|j>
      mur = - <i|r|a>
      muv = - <i|nabla|a>/omega_ia with omega_ia>0
      magn = <i|[r x nabla]|a> / (2 m_e c)
    """

    def __init__(self, iidx=None, jidx=None, pspin=None, kpt=None,
                 paw=None, string=None, fijscale=1, dtype=float):

        if string is not None:
            self.fromstring(string, dtype)
            return None

        # normal entry

        PairDensity.__init__(self, paw)
        PairDensity.initialize(self, kpt, iidx, jidx)

        self.pspin = pspin

        self.energy = 0.0
        self.fij = 0.0

        self.me = np.zeros((3), dtype=dtype)
        self.mur = np.zeros((3), dtype=dtype)
        self.muv = np.zeros((3), dtype=dtype)
        self.magn = np.zeros((3), dtype=dtype)

        self.kpt_comm = paw.wfs.kd.comm

        # leave empty if not my kpt
        if kpt is None:
            return

        wfs = paw.wfs
        gd = wfs.gd

        self.energy = kpt.eps_n[jidx] - kpt.eps_n[iidx]
        self.fij = (kpt.f_n[iidx] - kpt.f_n[jidx]) * fijscale

        # calculate matrix elements -----------

        # length form ..........................

        # course grid contribution
        # <i|r|j> is the negative of the dipole moment (because of negative
        # e- charge)
        me = - gd.calculate_dipole_moment(self.get())

        # augmentation contributions
        ma = np.zeros(me.shape, dtype=dtype)
        pos_av = paw.atoms.get_positions() / Bohr
        for a, P_ni in kpt.P_ani.items():
            Ra = pos_av[a]
            Pi_i = P_ni[self.i].conj()
            Pj_i = P_ni[self.j]
            Delta_pL = wfs.setups[a].Delta_pL
            ni = len(Pi_i)
            ma0 = 0
            ma1 = np.zeros(me.shape, dtype=me.dtype)
            for i in range(ni):
                for j in range(ni):
                    pij = Pi_i[i] * Pj_i[j]
                    ij = packed_index(i, j, ni)
                    # L=0 term
                    ma0 += Delta_pL[ij, 0] * pij
                    # L=1 terms
                    if wfs.setups[a].lmax >= 1:
                        # see spherical_harmonics.py for
                        # L=1:y L=2:z; L=3:x
                        ma1 += np.array([Delta_pL[ij, 3], Delta_pL[ij, 1],
                                         Delta_pL[ij, 2]]) * pij
            ma += np.sqrt(4 * np.pi / 3) * ma1 + Ra * np.sqrt(4 * np.pi) * ma0
        gd.comm.sum(ma)

        self.me = np.sqrt(self.energy * self.fij) * (me + ma)
        self.mur = - (me + ma)

        # velocity form .............................

        if self.lcao:
            # XXX Velocity form not supported in LCAO
            return

        me = np.zeros(self.mur.shape, dtype=dtype)

        # get derivatives
        dtype = self.wfj.dtype
        dwfj_cg = gd.empty((3), dtype=dtype)
        if not hasattr(gd, 'ddr'):
            gd.ddr = [Gradient(gd, c, dtype=dtype).apply for c in range(3)]
        for c in range(3):
            gd.ddr[c](self.wfj, dwfj_cg[c], kpt.phase_cd)
            me[c] = gd.integrate(self.wfi.conj() * dwfj_cg[c])

        if 0:
            me2 = np.zeros(self.mur.shape)
            for c in range(3):
                gd.ddr[c](self.wfi, dwfj_cg[c], kpt.phase_cd)
                me2[c] = gd.integrate(self.wfj * dwfj_cg[c])
            print(me, -me2, me2 + me)

        # augmentation contributions
        ma = np.zeros(me.shape, dtype=me.dtype)
        for a, P_ni in kpt.P_ani.items():
            Pi_i = P_ni[self.i].conj()
            Pj_i = P_ni[self.j]
            nabla_iiv = paw.wfs.setups[a].nabla_iiv
            for c in range(3):
                for i1, Pi in enumerate(Pi_i):
                    for i2, Pj in enumerate(Pj_i):
                        ma[c] += Pi * Pj * nabla_iiv[i1, i2, c]
        gd.comm.sum(ma)

        self.muv = - (me + ma) / self.energy

        # magnetic transition dipole ................

        r_cg, r2_g = coordinates(gd)
        magn = np.zeros(me.shape, dtype=dtype)

        wfi_g = self.wfi.conj()
        for ci in range(3):
            cj = (ci + 1) % 3
            ck = (ci + 2) % 3
            magn[ci] = gd.integrate(wfi_g * r_cg[cj] * dwfj_cg[ck] -
                                    wfi_g * r_cg[ck] * dwfj_cg[cj])
        # augmentation contributions
        ma = np.zeros(magn.shape, dtype=magn.dtype)
        for a, P_ni in kpt.P_ani.items():
            Pi_i = P_ni[self.i].conj()
            Pj_i = P_ni[self.j]
            rnabla_iiv = paw.wfs.setups[a].rnabla_iiv
            for c in range(3):
                for i1, Pi in enumerate(Pi_i):
                    for i2, Pj in enumerate(Pj_i):
                        ma[c] += Pi * Pj * rnabla_iiv[i1, i2, c]
        gd.comm.sum(ma)

        self.magn = -alpha / 2. * (magn + ma)

    def distribute(self):
        """Distribute results to all cores."""
        self.spin = self.kpt_comm.sum(self.spin)
        self.pspin = self.kpt_comm.sum(self.pspin)
        self.k = self.kpt_comm.sum(self.k)
        self.weight = self.kpt_comm.sum(self.weight)
        self.energy = self.kpt_comm.sum(self.energy)
        self.fij = self.kpt_comm.sum(self.fij)

        self.kpt_comm.sum(self.me)
        self.kpt_comm.sum(self.mur)
        self.kpt_comm.sum(self.muv)
        self.kpt_comm.sum(self.magn)

    def __add__(self, other):
        """Add two KSSingles"""
        result = self.copy()
        result.me = self.me + other.me
        result.mur = self.mur + other.mur
        result.muv = self.muv + other.muv
        result.magn = self.magn - other.magn
        return result

    def __sub__(self, other):
        """Subtract two KSSingles"""
        result = self.copy()
        result.me = self.me - other.me
        result.mur = self.mur - other.mur
        result.muv = self.muv - other.muv
        result.magn = self.magn - other.magn
        return result

    def __rmul__(self, x):
        return self.__mul__(x)

    def __mul__(self, x):
        """Multiply a KSSingle with a number"""
        if isinstance(x, (float, int)):
            result = self.copy()
            result.me = self.me * x
            result.mur = self.mur * x
            result.muv = self.muv * x
            return result
        else:
            return RuntimeError('not a number')

    def __truediv__(self, x):
        return self.__mul__(1. / x)

    __div__ = __truediv__

    def copy(self):
        if self.mur.dtype == complex:
            return KSSingle(string=self.outstring(), dtype=complex)
        else:
            return KSSingle(string=self.outstring(), dtype=float)

    def fromstring(self, string, dtype=float):
        l = string.split()
        self.i = int(l.pop(0))
        self.j = int(l.pop(0))
        self.pspin = int(l.pop(0))
        self.spin = int(l.pop(0))
        if dtype == float:
            self.k = 0
            self.weight = 1
        else:
            self.k = int(l.pop(0))
            self.weight = float(l.pop(0))
        self.energy = float(l.pop(0))
        self.fij = float(l.pop(0))
        self.mur = np.array([dtype(l.pop(0)) for i in range(3)])
        self.me = - self.mur * np.sqrt(self.energy * self.fij)
        self.muv = self.magn = None
        if len(l):
            self.muv = np.array([dtype(l.pop(0)) for i in range(3)])
        if len(l):
            self.magn = np.array([dtype(l.pop(0)) for i in range(3)])
        return None

    def outstring(self):
        if self.mur.dtype == float:
            string = '{0:d} {1:d}  {2:d} {3:d}  {4:g} {5:g}'.format(
                self.i, self.j, self.pspin, self.spin, self.energy, self.fij)
        else:
            string = (
                '{0:d} {1:d}  {2:d} {3:d} {4:d} {5:g}  {6:g} {7:g}'.format(
                    self.i, self.j, self.pspin, self.spin, self.k,
                    self.weight, self.energy, self.fij))
        string += '  '

        def format_me(me):
            string = ''
            if me.dtype == float:
                for m in me:
                    string += ' {0:.5e}'.format(m)
            else:
                for m in me:
                    string += ' {0.real:.5e}{0.imag:+.5e}j'.format(m)
            return string

        string += '  ' + format_me(self.mur)
        if self.muv is not None:
            string += '  ' + format_me(self.muv)
        if self.magn is not None:
            string += '  ' + format_me(self.magn)
        string += '\n'

        return string

    def __str__(self):
        string = '# <KSSingle> %d->%d %d(%d) eji=%g[eV]' % \
            (self.i, self.j, self.pspin, self.spin,
             self.energy * Hartree)
        if self.me.dtype == float:
            string += ' (%g,%g,%g)' % (self.me[0], self.me[1], self.me[2])
        else:
            string += ' kpt={0:d} w={1:g}'.format(self.k, self.weight)
            string += ' ('
            # use velocity form
            s = - np.sqrt(self.energy * self.fij)
            for c, m in enumerate(s * self.me):
                string += '{0.real:.5e}{0.imag:+.5e}j'.format(m)
                if c < 2:
                    string += ','
            string += ')'
        return string

    def __eq__(self, other):
        """KSSingles are considred equal when their indices are equal."""
        return (self.pspin == other.pspin and self.k == other.k and
                self.i == other.i and self.j == other.j)

    def __hash__(self):
        """Hash similar to __eq__"""
        if not hasattr(self, 'hash'):
            self.hash = hash((self.spin, self.k, self.i, self.j))
        return self.hash

    #
    # User interface: ##
    #

    def get_weight(self):
        return self.fij
