"""This module is used to generate atomic orbital basis sets."""
import sys
from io import StringIO

import numpy as np
from ase.units import Hartree
from gpaw.utilities import devnull

from gpaw import __version__ as version
from gpaw.utilities import divrl
from gpaw.atom.generator import Generator
from gpaw.atom.all_electron import AllElectron
from gpaw.atom.configurations import parameters
from gpaw.basis_data import Basis, BasisFunction, get_basis_name
from gpaw.atom.radialgd import (AERadialGridDescriptor,
                                EquidistantRadialGridDescriptor)


def get_basis_l(f_j, l_j, Nv):
    N = 0
    lvalues = []
    j_l = []
    Nl_j = len(l_j)
    while N < Nv:
        Nl_j -= 1
        if f_j[Nl_j] > 0:
            lvalues.append(l_j[Nl_j])
            N += f_j[Nl_j]
            j_l.append(Nl_j)
    assert N == Nv

    if min(lvalues) != 0:  # always include s-orbital !
        lvalues = lvalues + [0]
        reversed_l_j = list(l_j)
        reversed_l_j.reverse()
        j = len(reversed_l_j) - reversed_l_j.index(0) - 1
        j_l.append(j)

    return [j_l, lvalues]


def get_gaussianlike_basis_function(rgd, l, rchar, gcut):
    rcut = rgd.r_g[gcut]
    gaussian = QuasiGaussian(1.0 / rchar**2, rcut)
    r_g = rgd.r_g[:gcut + 1]
    phit_g = gaussian(r_g) * r_g**l
    phit_g[-1] = 0.0
    norm = np.dot(rgd.dr_g[:gcut + 1], (r_g * phit_g)**2)**0.5
    phit_g /= norm
    norm2 = np.dot(rgd.dr_g[:gcut + 1], (r_g * phit_g)**2)**0.5
    assert abs(1.0 - norm2) < 1e-10
    return phit_g


def make_split_valence_basis_function(r_g, psi_g, l, gcut):
    """Get polynomial which joins psi smoothly at rcut.

    Returns an array of function values f(r) * r, where::

              l           2
      f(r) = r  * (a - b r ),  r < rcut
      f(r) = psi(r),           r >= rcut

    where a and b are determined such that f(r) is continuous and
    differentiable at rcut.  The parameter psi should be an atomic
    orbital.
    """
    r1 = r_g[gcut]  # ensure that rcut is moved to a grid point
    r2 = r_g[gcut + 1]
    y1 = psi_g[gcut] / r_g[gcut]
    y2 = psi_g[gcut + 1] / r_g[gcut + 1]
    b = - (y2 / r2**l - y1 / r1**l) / (r2**2 - r1**2)
    a = (y1 / r1**l + b * r1**2)
    psi_g2 = r_g**(l + 1) * (a - b * r_g**2)
    psi_g2[gcut:] = psi_g[gcut:]
    return psi_g2


def rsplit_by_norm(rgd, l, u, tailnorm_squared, txt):
    """Find radius outside which remaining tail has a particular norm."""
    norm_squared = np.dot(rgd.dr_g, u * u)
    partial_norm_squared = 0.
    i = len(u) - 1
    absolute_tailnorm_squared = tailnorm_squared * norm_squared
    while partial_norm_squared < absolute_tailnorm_squared:
        # Integrate backwards.  This is important since the pseudo
        # wave functions have strange behaviour near the core.
        partial_norm_squared += rgd.dr_g[i] * u[i]**2
        i -= 1
    rsplit = rgd.r_g[i + 1]
    msg = ('Tail norm %.03f :: rsplit=%.02f Bohr' %
           ((partial_norm_squared / norm_squared)**0.5, rsplit))
    print(msg, file=txt)
    gsplit = rgd.floor(rsplit)
    splitwave = make_split_valence_basis_function(rgd.r_g, u, l, gsplit)
    return rsplit, partial_norm_squared, splitwave


class BasisMaker:
    """Class for creating atomic basis functions."""
    def __init__(self, generator, name=None, run=True, gtxt='-',
                 non_relativistic_guess=False, xc='PBE',
                 save_setup=False):

        if isinstance(generator, str):  # treat 'generator' as symbol
            generator = Generator(generator, scalarrel=True,
                                  xcname=xc, txt=gtxt,
                                  nofiles=True)
            generator.N *= 4
        self.generator = generator
        self.rgd = AERadialGridDescriptor(generator.beta / generator.N,
                                          1.0 / generator.N, generator.N,
                                          default_spline_points=100)
        self.name = name

        if run:
            if non_relativistic_guess:
                ae0 = AllElectron(generator.symbol, scalarrel=False,
                                  nofiles=False, txt=gtxt, xcname=xc)
                ae0.N = generator.N
                ae0.beta = generator.beta
                ae0.run()
                # Now files will be stored such that they can
                # automagically be used by the next run()
            setup = generator.run(write_xml=False,
                                  name=name,
                                  **parameters[generator.symbol])

            if save_setup:
                setup.write_xml()
        else:
            if save_setup:
                raise ValueError('cannot save setup here because setup '
                                 'was already generated before basis '
                                 'generation.')

    def smoothify(self, psi_mg, l):
        r"""Generate pseudo wave functions from all-electron ones.

        The pseudo wave function is::

                                   ___
               ~                   \    /   ~             \    ~     ~
            | psi  > = | psi  > +   )  | | phi > - | phi > ) < p  | psi > ,
                 m          m      /__  \     i         i /     i      m
                                    i

        where the scalar products are found by solving::

                            ___
              ~             \     ~             ~     ~
            < p | psi  > =   )  < p  | phi  > < p  | psi  > .
               i     m      /__    i      j      j      m
                             j

        In order to ensure smoothness close to the core, the
        all-electron wave function and partial wave are then
        multiplied by a radial function which approaches 0 near the
        core, such that the pseudo wave function approaches::

                        ___
               ~        \      ~       ~     ~
            | psi  > =   )  | phi >  < p  | psi >    (for r << rcut),
                 m      /__      i      i      m
                         i

        which is exact if the projectors/pseudo partial waves are complete.
        """
        if psi_mg.ndim == 1:
            return self.smoothify(psi_mg[None], l)[0]

        g = self.generator
        u_ng = g.u_ln[l]
        q_ng = g.q_ln[l]
        s_ng = g.s_ln[l]

        Pi_nn = np.dot(g.r * q_ng, u_ng.T)
        Q_nm = np.dot(g.r * q_ng, psi_mg.T)
        Qt_nm = np.linalg.solve(Pi_nn, Q_nm)

        # Weight-function for truncating all-electron parts smoothly near core
        gmerge = g.r2g(g.rcut_l[l])
        w_g = np.ones(g.r.shape)
        w_g[0:gmerge] = (g.r[0:gmerge] / g.r[gmerge])**2.
        w_g = w_g[None]

        psit_mg = psi_mg * w_g + np.dot(Qt_nm.T, s_ng - u_ng * w_g)
        return psit_mg

    def make_orbital_vector(self, j, rcut, vconf=None):
        """Returns a smooth basis vector given an all-electron one."""
        l = self.generator.l_j[j]
        psi_g, e = self.generator.solve_confined(j, rcut, vconf)
        psit_g = self.smoothify(psi_g, l)
        return psit_g

    def rcut_by_energy(self, j, esplit=.1, tolerance=.1, rguess=6.,
                       vconf_args=None):
        """Find confinement cutoff corresponding to given orbital energy shift.

        Creates a confinement potential for the orbital given by j,
        such that the confined-orbital energy is (emin to emax) eV larger
        than the free-orbital energy."""
        g = self.generator
        e_base = g.e_j[j]
        rc = rguess

        if vconf_args is None:
            vconf = None
        else:
            amplitude, ri_rel = vconf_args
            vconf = g.get_confinement_potential(amplitude, ri_rel * rc, rc)

        psi_g, e = g.solve_confined(j, rc, vconf)
        de_min, de_max = esplit / Hartree, (esplit + tolerance) / Hartree

        rmin = 0.
        rmax = g.r[-1]

        de = e - e_base
        # print '--------'
        # print 'Start bisection'
        # print'e_base =',e_base
        # print 'e =',e
        # print '--------'
        while de < de_min or de > de_max:
            if de < de_min:  # Move rc left -> smaller cutoff, higher energy
                rmax = rc
                rc = (rc + rmin) / 2.
            else:  # Move rc right
                rmin = rc
                rc = (rc + rmax) / 2.
            if vconf is not None:
                vconf = g.get_confinement_potential(amplitude, ri_rel * rc, rc)
            psi_g, e = g.solve_confined(j, rc, vconf)
            de = e - e_base
            # print 'rc = %.03f :: e = %.03f :: de = %.03f' % (rc, e*Hartree,
            #                                                  de*Hartree)
            # if rmin - rmax < 1e-
            if g.r2g(rmax) - g.r2g(rmin) <= 1:  # adjacent points
                break  # cannot meet tolerance due to grid resolution
        # print 'Done!'
        return psi_g, e, de, vconf, rc

    def generate(self, zetacount=2, polarizationcount=1,
                 tailnorm=(0.16, 0.3, 0.6), energysplit=0.1, tolerance=1.0e-3,
                 referencefile=None, referenceindex=None, rcutpol_rel=1.0,
                 rcutmax=20.0,
                 rcharpol_rel=None,
                 vconf_args=(12.0, 0.6), txt='-',
                 include_energy_derivatives=False,
                 # lvalues=None, # XXX clean up some of these!
                 jvalues=None,
                 l_pol=None):
        """Generate an entire basis set.

        This is a high-level method which will return a basis set
        consisting of several different basis vector types.

        Parameters:

        ===================== =================================================
        ``zetacount``         Number of basis functions per occupied orbital
        ``polarizationcount`` Number of polarization functions
        ``tailnorm``          List of tail norms for split-valence scheme
        ``energysplit``       Energy increase defining confinement radius (eV)
        ``tolerance``         Tolerance of energy split (eV)
        ``referencefile``     gpw-file used to generate polarization function
        ``referenceindex``    Index in reference system of relevant atom
        ``rcutpol_rel``       Polarization rcut relative to largest other rcut
        ``rcutmax``           No cutoff will be greater than this value
        ``vconf_args``        Parameters (alpha, ri/rc) for conf. potential
        ``txt``               Log filename or '-' for stdout
        ===================== =================================================

        Returns a fully initialized Basis object.
        """

        if txt == '-':
            txt = sys.stdout
        elif txt is None:
            txt = devnull

        if isinstance(tailnorm, float):
            tailnorm = (tailnorm,)
        if 1 + len(tailnorm) < max(polarizationcount, zetacount):
            raise ValueError(
                'Needs %d tail norm values, but only %d are specified' %
                (max(polarizationcount, zetacount) - 1, len(tailnorm)))

        textbuffer = StringIO()

        class TeeStream:  # quick hack to both write and save output
            def __init__(self, out1, out2):
                self.out1 = out1
                self.out2 = out2

            def write(self, string):
                self.out1.write(string)
                self.out2.write(string)

        txt = TeeStream(txt, textbuffer)

        if vconf_args is not None:
            amplitude, ri_rel = vconf_args

        g = self.generator
        rgd = self.rgd

        njcore = g.njcore
        n_j = g.n_j[njcore:]
        l_j = g.l_j[njcore:]
        f_j = g.f_j[njcore:]

        if jvalues is None:
            jvalues = []
            sortkeys = []
            for j in range(len(n_j)):
                if f_j[j] == 0 and l_j[j] != 0:
                    continue
                jvalues.append(j)
                sortkeys.append(l_j[j])

            # Now order jvalues by l
            #
            # Use a stable sort so the energy ordering within each
            # angular momentum is guaranteed to be preserved
            args = np.argsort(sortkeys, kind='mergesort')
            jvalues = np.array(jvalues)[args]

        fulljvalues = [njcore + j for j in jvalues]

        if isinstance(energysplit, float):
            energysplit = [energysplit] * len(jvalues)

        title = '%s Basis functions for %s' % (g.xcname, g.symbol)
        print(title, file=txt)
        print('=' * len(title), file=txt)

        singlezetas = []
        energy_derivative_functions = []
        multizetas = [[] for i in range(zetacount - 1)]
        polarization_functions = []

        splitvalencedescr = 'split-valence wave, fixed tail norm'
        derivativedescr = 'derivative of sz wrt. (ri/rc) of potential'

        for vj, fullj, esplit in zip(jvalues, fulljvalues, energysplit):
            l = l_j[vj]
            n = n_j[vj]
            assert n > 0
            orbitaltype = str(n) + 'spdf'[l]
            msg = 'Basis functions for l=%d, n=%d' % (l, n)
            print(file=txt)
            print(msg + '\n', '-' * len(msg), file=txt)
            print(file=txt)
            if vconf_args is None:
                adverb = 'sharply'
            else:
                adverb = 'softly'
            print('Zeta 1: %s confined pseudo wave,' % adverb, end=' ',
                  file=txt)

            u, e, de, vconf, rc = self.rcut_by_energy(fullj, esplit,
                                                      tolerance,
                                                      vconf_args=vconf_args)
            if rc > rcutmax:
                rc = rcutmax  # scale things down
                if vconf is not None:
                    vconf = g.get_confinement_potential(amplitude, ri_rel * rc,
                                                        rc)
                u, e = g.solve_confined(fullj, rc, vconf)
                print('using maximum cutoff', file=txt)
                print('rc=%.02f Bohr' % rc, file=txt)
            else:
                print('fixed energy shift', file=txt)
                print('DE=%.03f eV :: rc=%.02f Bohr'
                      % (de * Hartree, rc), file=txt)

            if vconf is not None:
                print('Potential amp=%.02f :: ri/rc=%.02f' %
                      (amplitude, ri_rel), file=txt)
            phit_g = self.smoothify(u, l)
            bf = BasisFunction(n, l, rc, phit_g,
                               '%s-sz confined orbital' % orbitaltype)
            norm = np.dot(g.dr, phit_g * phit_g)**.5
            print('Norm=%.03f' % norm, file=txt)
            singlezetas.append(bf)

            zetacounter = iter(range(2, zetacount + 1))

            if include_energy_derivatives:
                assert zetacount > 1
                zeta = next(zetacounter)
                print('\nZeta %d: %s' % (zeta, derivativedescr), file=txt)
                vconf2 = g.get_confinement_potential(amplitude,
                                                     ri_rel * rc * .99, rc)
                u2, e2 = g.solve_confined(fullj, rc, vconf2)

                phit2_g = self.smoothify(u2, l)
                dphit_g = phit2_g - phit_g
                dphit_norm = np.dot(rgd.dr_g, dphit_g * dphit_g) ** .5
                dphit_g /= dphit_norm
                descr = '%s-dz E-derivative of sz' % orbitaltype
                bf = BasisFunction(None, l, rc, dphit_g, descr)
                energy_derivative_functions.append(bf)

            for i, zeta in enumerate(zetacounter):
                print('\nZeta %d: %s' % (zeta, splitvalencedescr), file=txt)
                # Unresolved issue:  how does the lack of normalization
                # of the first function impact the tail norm scheme?
                # Presumably not much, since most interesting stuff happens
                # close to the core.
                rsplit, norm, splitwave = rsplit_by_norm(rgd, l, phit_g,
                                                         tailnorm[i]**2.0,
                                                         txt)
                descr = '%s-%sz split-valence wave' % (orbitaltype,
                                                       '0sdtq56789'[zeta])
                bf = BasisFunction(None, l, rsplit, phit_g - splitwave, descr)
                multizetas[i].append(bf)

        if polarizationcount > 0 or l_pol is not None:
            if l_pol is None:
                # Now make up some properties for the polarization orbital
                # We just use the cutoffs from the previous one times a factor
                # Find 'missing' values in lvalues
                lvalues = [l_j[vj] for vj in jvalues]
                for i in range(max(lvalues) + 1):
                    if list(lvalues).count(i) == 0:
                        l_pol = i
                        break
                else:
                    l_pol = max(lvalues) + 1

            # Find the last state with l=l_pol - 1, which will be the state we
            # base the polarization function on
            for vj, fullj, bf in zip(jvalues[::-1], fulljvalues[::-1],
                                     singlezetas[::-1]):
                if bf.l == l_pol - 1:
                    fullj_pol = fullj
                    rcut = bf.rc * rcutpol_rel
                    break
            else:
                raise ValueError('The requested value l_pol=%d requires l=%d '
                                 'among valence states' % (l_pol, l_pol - 1))
            rcut = min(rcut, rcutmax)
            msg = 'Polarization function: l=%d, rc=%.02f' % (l_pol, rcut)
            print('\n' + msg, file=txt)
            print('-' * len(msg), file=txt)
            # Make a single Gaussian for polarization function.
            #
            # It is known that for given l, the sz cutoff defined
            # by some fixed energy is strongly correlated to the
            # value of the characteristic radius which best reproduces
            # the wave function found by interpolation.
            #
            # We know that for e.g. d orbitals:
            #   rchar ~= .37 rcut[sz](.3eV)
            # Since we don't want to spend a lot of time finding
            # these value for other energies, we just find the energy
            # shift at .3 eV now

            u, e, de, vconf, rc_fixed = self.rcut_by_energy(fullj_pol,
                                                            .3, 1e-2,
                                                            6., (12., .6))

            default_rchar_rel = .25
            # Defaults for each l.  Actually we don't care right now
            rchar_rels = {}

            if rcharpol_rel is None:
                rcharpol_rel = rchar_rels.get(l_pol, default_rchar_rel)
            rchar = rcharpol_rel * rc_fixed
            gaussian = QuasiGaussian(1.0 / rchar**2, rcut)
            psi_pol = gaussian(rgd.r_g) * rgd.r_g**(l_pol + 1)
            norm = np.dot(rgd.dr_g, psi_pol * psi_pol) ** .5
            psi_pol /= norm
            print('Single quasi Gaussian', file=txt)
            msg = 'Rchar = %.03f*rcut = %.03f Bohr' % (rcharpol_rel, rchar)
            adjective = 'Gaussian'
            print(msg, file=txt)
            type = '%s-type %s polarization' % ('spdfg'[l_pol], adjective)
            bf_pol = BasisFunction(None, l_pol, rcut, psi_pol, type)

            polarization_functions.append(bf_pol)
            for i in range(polarizationcount - 1):
                npol = i + 2
                msg = '\n%s: %s' % (['Secondary', 'Tertiary', 'Quaternary',
                                     'Quintary', 'Sextary', 'Septenary'][i],
                                    splitvalencedescr)
                print(msg, file=txt)
                rsplit, norm, splitwave = rsplit_by_norm(rgd, l_pol, psi_pol,
                                                         tailnorm[i],
                                                         txt)
                descr = ('%s-type split-valence polarization %d'
                         % ('spdfg'[l_pol], npol))
                bf_pol = BasisFunction(None, l_pol, rsplit,
                                       psi_pol - splitwave,
                                       descr)
                polarization_functions.append(bf_pol)

        bf_j = []
        bf_j.extend(singlezetas)
        bf_j.extend(energy_derivative_functions)
        for multizeta_list in multizetas:
            bf_j.extend(multizeta_list)
        bf_j.extend(polarization_functions)

        rcmax = max([bf.rc for bf in bf_j])

        # The non-equidistant grids are really only suited for AE WFs
        d = 1.0 / 64
        equidistant_grid = np.arange(0.0, rcmax + d, d)
        ng = len(equidistant_grid)

        for bf in bf_j:
            # We have been storing phit_g * r, but we just want phit_g
            bf.phit_g = divrl(bf.phit_g, 1, rgd.r_g)

            gcut = min(int(1 + bf.rc / d), ng - 1)

            assert equidistant_grid[gcut] >= bf.rc
            assert equidistant_grid[gcut - 1] <= bf.rc

            bf.rc = equidistant_grid[gcut]
            # Note: bf.rc *must* correspond to a grid point (spline issues)
            bf.ng = gcut + 1
            # XXX all this should be done while building the basis vectors,
            # not here

            # Quick hack to change to equidistant coordinates
            spline = rgd.spline(bf.phit_g, rgd.r_g[rgd.floor(bf.rc)], bf.l,
                                points=100)
            bf.phit_g = np.array([spline(r) * r**bf.l
                                  for r in equidistant_grid[:bf.ng]])
            bf.phit_g[-1] = 0.

        basistype = get_basis_name(zetacount, polarizationcount)
        if self.name is None:
            compound_name = basistype
        else:
            compound_name = '%s.%s' % (self.name, basistype)

        basis = Basis(g.symbol, compound_name, False,
                      EquidistantRadialGridDescriptor(d, ng))
        basis.bf_j = bf_j
        basis.generatordata = textbuffer.getvalue().strip()
        basis.generatorattrs = {'version': version}
        textbuffer.close()

        return basis

    def grplot(self, bf_j):
        """Plot basis functions on generator's radial grid."""
        import matplotlib.pyplot as plt
        rc = max([bf.rc for bf in bf_j])
        g = self.generator
        r = g.r
        for bf in bf_j:
            label = bf.type
            # XXX times g.r or not times g.r ?
            plt.plot(r, bf.phit_g / r, label=label[:12])
        axis = plt.axis()
        newaxis = [0., rc, axis[2], axis[3]]
        plt.axis(newaxis)
        plt.legend()
        plt.show()

    def plot(self, basis, figure=None, title=None, filename=None):
        """Plot basis functions using MatPlotLib."""
        # XXX method should no longer belong to a basis maker
        import matplotlib.pyplot as plt
        rc = max([bf.rc for bf in basis.bf_j])
        r = np.linspace(0., basis.d * (basis.ng - 1), basis.ng)
        g = self.generator
        if figure is not None:
            plt.figure(figure)
        else:
            plt.figure()  # not very elegant
        if title is None:
            title = g.symbol
        plt.title(title)
        for bf in basis.bf_j:
            label = bf.type
            # XXX times g.r or not times g.r ?
            phit_g = np.zeros_like(r)
            phit_g[:len(bf.phit_g)] = bf.phit_g
            plt.plot(r, phit_g * r, label=label[:12])
        axis = plt.axis()
        newaxis = [0., rc, axis[2], axis[3]]
        plt.axis(newaxis)
        plt.legend()
        if filename is not None:
            plt.savefig(filename)


class QuasiGaussian:
    """Gaussian-like functions for expansion of orbitals.

    Implements f(r) = A [G(r) - P(r)] where::

      G(r) = exp{- alpha r^2}
      P(r) = a - b r^2

    with (a, b) such that f(rcut) == f'(rcut) == 0.
    """
    def __init__(self, alpha, rcut, A=1.):
        self.alpha = alpha
        self.rcut = rcut
        expmar2 = np.exp(-alpha * rcut**2)
        a = (1 + alpha * rcut**2) * expmar2
        b = alpha * expmar2
        self.a = a
        self.b = b
        self.A = A

    def __call__(self, r):
        """Evaluate function values at r, which is a numpy array."""
        condition = (r < self.rcut) & (self.alpha * r**2 < 700.)
        r2 = np.where(condition, r**2., 0.)  # prevent overflow
        g = np.exp(-self.alpha * r2)
        p = (self.a - self.b * r2)
        y = np.where(condition, g - p, 0.)
        return self.A * y

    def renormalize(self, norm):
        """Divide function by norm."""
        self.A /= norm
