"""This module is used to generate atomic orbital basis sets."""
from __future__ import print_function
import sys
try:
    from StringIO import StringIO  # Python 2
except ImportError:
    from io import StringIO

import numpy as np
from ase.units import Hartree
from ase.utils 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, use_restart_file=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 pylab
        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 ?
            pylab.plot(r, bf.phit_g / r, label=label[:12])
        axis = pylab.axis()
        newaxis = [0., rc, axis[2], axis[3]]
        pylab.axis(newaxis)
        pylab.legend()
        pylab.show()

    def plot(self, basis, figure=None, title=None, filename=None):
        """Plot basis functions using pylab."""
        # XXX method should no longer belong to a basis maker
        import pylab
        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:
            pylab.figure(figure)
        else:
            pylab.figure()  # not very elegant
        if title is None:
            title = g.symbol
        pylab.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
            pylab.plot(r, phit_g * r, label=label[:12])
        axis = pylab.axis()
        newaxis = [0., rc, axis[2], axis[3]]
        pylab.axis(newaxis)
        pylab.legend()
        if filename is not None:
            pylab.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
