import numbers
from abc import ABC, abstractmethod
from math import factorial as fac
from math import pi
from typing import Tuple

import numpy as np
from gpaw.sphere.integrate import integrate_radial_grid
from gpaw.spline import Spline
from gpaw.typing import Array1D
from gpaw.utilities import divrl, hartree
from scipy.interpolate import make_interp_spline, splder


def radial_grid_descriptor(eq: str, **kwargs) -> 'RadialGridDescriptor':
    if eq == 'r=d*i':
        assert int(kwargs['istart']) == 0
        return EquidistantRadialGridDescriptor(float(kwargs['d']),
                                               int(kwargs['iend']) + 1)
    if eq == 'r=a*i/(n-i)':
        beta = float(kwargs['a'])
        ng = int(kwargs['n'])
        return AERadialGridDescriptor(beta / ng, 1.0 / ng, ng)
    if eq == 'r=a*i/(1-b*i)':
        a = float(kwargs['a'])
        b = float(kwargs['b'])
        N = int(kwargs['n'])
        return AERadialGridDescriptor(a, b, N)
    if eq == 'r=a*(exp(d*i)-1)':
        a = float(kwargs['a'])
        d = float(kwargs['d'])
        N = int(kwargs['iend']) + 1
        return AbinitRadialGridDescriptor(a, d, N)
    raise ValueError('Unknown grid: ' + eq)


def fsbt(l, f_g, r_g, G_k):
    """Fast spherical Bessel transform.

    Returns::

          oo
         / 2
         |r dr j (Gr) f(r),
         /      l
          0

    using l+1 fft's."""

    N = (len(G_k) - 1) * 2
    f_k = 0.0
    F_g = f_g * r_g
    for n in range(l + 1):
        f_k += (r_g[1] * (1j)**(l + 1 - n) *
                fac(l + n) / fac(l - n) / fac(n) / 2**n *
                np.fft.rfft(F_g, N)).real * G_k**(l - n)
        F_g[1:] /= r_g[1:]

    f_k[1:] /= G_k[1:]**(l + 1)
    if l == 0:
        f_k[0] = np.dot(r_g, f_g * r_g) * r_g[1]
    return f_k


class RadialGridDescriptor(ABC):
    ndim = 1  # dimension of ndarrays

    def __init__(self, r_g: np.ndarray, dr_g, default_spline_points=25):
        """Grid descriptor for radial grid."""
        self.r_g = r_g
        self.dr_g = dr_g
        self.N = len(r_g)
        self.dv_g = 4 * pi * r_g**2 * dr_g
        self.default_spline_points = default_spline_points

    def __len__(self):
        return self.N

    @classmethod
    def new(cls, name, *args, **kwargs):
        if name == 'equidistant':
            return EquidistantRadialGridDescriptor(*args, **kwargs)
        raise ValueError(f'Unknown grid-type: {name}')

    def zeros(self, x=()):
        a_xg = self.empty(x)
        a_xg[:] = 0
        return a_xg

    def empty(self, x=()):
        if isinstance(x, int):
            x = (x,)
        return np.zeros(x + (self.N,))

    def integrate(self, a_xg, n=0):
        assert n >= -2
        return np.dot(a_xg[..., 1:],
                      (self.r_g**(2 + n) * self.dr_g)[1:]) * (4 * pi)

    def integrate_trapz(self, a_xg, rcut=None):
        return integrate_radial_grid(a_xg, self.r_g, rcut=rcut)

    def yukawa(self, n_g, l=0, gamma=1e-6):
        r"""Calculates the radial grid yukawa integral.

        The the integral kernel for the Yukawa interaction:

                    \    _   _
              exp(- /\ | r - r' |)
              ----------------------
                      _   _
                    | r - r' |

           is defined as

            __    __            \  r              \  r    * ^     ^
          \     4 ||  I_(l+0.5)(/\  <) K_(l+0.5) (/\  >) Y (r)  Y(r')
           )          --------------------------          lm     lm
          / __            (rr')^0.5
            lm

         where I and K are the modified Bessel functions of the first
         and second kind (K is also known as Macdonald function).
         r = min (r, r')     r = max(r, r')
          <                   >
         We now calculate the integral:


                  ^    / _           ^
         v (r) Y (r) = |dr' n(r') Y (r')
          l     lm     /     l     lm

        with the Yukawa kernel mentioned above.

        And the output array is 'vr' as it is
        within the Hartree / radial Poisson solver.
        """

        from scipy.special import iv, kv
        vr_g = self.zeros()
        nrdr_g = n_g * self.r_g**1.5 * self.dr_g
        p = 0
        q = 0
        k_rgamma = kv(l + 0.5, self.r_g * gamma)      # K(>)
        i_rgamma = iv(l + 0.5, self.r_g * gamma)      # I(<)
        k_rgamma[0] = kv(l + 0.5, self.r_g[1] * gamma * 1e-5)
        # We have two integrals: one for r< and one for r>
        # This loop-technique helps calculate them in once
        for g_ind in range(len(nrdr_g) - 1, -1, -1):
            dp = k_rgamma[g_ind] * nrdr_g[g_ind]  # r' is r>
            dq = i_rgamma[g_ind] * nrdr_g[g_ind]  # r' is r<
            vr_g[g_ind] = (p + 0.5 * dp) * i_rgamma[g_ind] - \
                          (q + 0.5 * dq) * k_rgamma[g_ind]
            p += dp
            q += dq
        vr_g[:] += q * k_rgamma[:]
        vr_g *= 4 * pi
        vr_g[:] *= self.r_g[:]**0.5
        return vr_g

    def derivative_spline(self, n_g, dndr_g=None, *, degree=5):
        """Spline-based derivative of radial function."""
        if dndr_g is None:
            dndr_g = self.empty()

        # Boundary conditions
        assert degree in [3, 5]
        if degree == 5:
            bc_type = ([(2, 0.0), (4, 0.0)], [(2, 0.0), (4, 0.0)])
        elif degree == 3:
            bc_type = ([(2, 0.0)], [(2, 0.0)])

        s = make_interp_spline(self.r_g, n_g, k=degree, bc_type=bc_type)
        s = splder(s)
        dndr_g[:] = s(self.r_g)
        return dndr_g

    def derivative(self, n_g, dndr_g=None):
        """Finite-difference derivative of radial function."""
        if dndr_g is None:
            dndr_g = self.empty()
        dndr_g[0] = n_g[1] - n_g[0]
        dndr_g[1:-1] = 0.5 * (n_g[2:] - n_g[:-2])
        dndr_g[-1] = n_g[-1] - n_g[-2]
        dndr_g /= self.dr_g
        return dndr_g

    def derivative2(self, a_g, b_g):
        """Finite-difference derivative of radial function.

        For an infinitely dense grid, this method would be identical
        to the `derivative` method."""

        c_g = a_g / self.dr_g
        b_g[0] = 0.5 * c_g[1] + c_g[0]
        b_g[1] = 0.5 * c_g[2] - c_g[0]
        b_g[1:-1] = 0.5 * (c_g[2:] - c_g[:-2])
        b_g[-2] = c_g[-1] - 0.5 * c_g[-3]
        b_g[-1] = -c_g[-1] - 0.5 * c_g[-2]

    def laplace(self, n_g, d2ndr2_g=None):
        """Laplace of radial function."""
        if d2ndr2_g is None:
            d2ndr2_g = self.empty()
        dndg_g = 0.5 * (n_g[2:] - n_g[:-2])
        d2ndg2_g = n_g[2:] - 2 * n_g[1:-1] + n_g[:-2]
        d2ndr2_g[1:-1] = (d2ndg2_g / self.dr_g[1:-1]**2 +
                          dndg_g * (self.d2gdr2()[1:-1] +
                                    2 / self.r_g[1:-1] / self.dr_g[1:-1]))
        d2ndr2_g[0] = d2ndr2_g[1]
        d2ndr2_g[-1] = d2ndr2_g[-2]
        return d2ndr2_g

    def T(self, u_g, l):
        dudg_g = 0.5 * (u_g[2:] - u_g[:-2])
        d2udg2_g = u_g[2:] - 2 * u_g[1:-1] + u_g[:-2]
        Tu_g = self.empty()
        Tu_g[1:-1] = -0.5 * (d2udg2_g / self.dr_g[1:-1]**2 +
                             dudg_g * self.d2gdr2()[1:-1])
        Tu_g[-1] = Tu_g[-2]
        Tu_g[1:] += 0.5 * l * (l + 1) * u_g[1:] / self.r_g[1:]**2
        Tu_g[0] = Tu_g[1]
        return Tu_g

    def interpolate(self, f_g, r_x):
        from scipy.interpolate import InterpolatedUnivariateSpline
        return InterpolatedUnivariateSpline(self.r_g, f_g)(r_x)

    def fft(self, fr_g, l=0, N=None):
        """Fourier transform.

        Returns G and f(G) arrays::

                                          _ _
               l    ^    / _         ^   iG.r
          f(G)i Y  (G) = |dr f(r)Y  (r) e    .
                 lm      /        lm
        """

        if N is None:
            N = 2**13

        assert N % 2 == 0

        r_x = np.linspace(0, self.r_g[-1], N)
        f_x = self.interpolate(fr_g, r_x)
        f_x[1:] /= r_x[1:]
        f_x[0] = f_x[1]
        G_k = np.linspace(0, pi / r_x[1], N // 2 + 1)
        f_k = 4 * pi * fsbt(l, f_x, r_x, G_k)
        return G_k, f_k

    def filter(self, f_g, rcut, Gcut, l=0):
        Rcut = 100.0
        N = 1024 * 8
        r_x = np.linspace(0, Rcut, N, endpoint=False)
        h = Rcut / N

        alpha = 4.0 / rcut**2
        mcut = np.exp(-alpha * rcut**2)
        r2_x = r_x**2
        m_x = np.exp(-alpha * r2_x)
        for n in range(2):
            m_x -= (alpha * (rcut**2 - r2_x))**n * (mcut / fac(n))
        xcut = int(np.ceil(rcut / r_x[1]))
        m_x[xcut:] = 0.0

        G_k = np.linspace(0, pi / h, N // 2 + 1)

        # Zeropad the function to same length as coordinates:
        fpad_g = np.zeros(len(self.r_g))
        fpad_g[:len(f_g)] = f_g
        f_g = fpad_g

        from scipy.interpolate import InterpolatedUnivariateSpline
        if l < 2:
            f_x = InterpolatedUnivariateSpline(self.r_g, f_g)(r_x)
        else:
            a_g = f_g.copy()
            a_g[1:] /= self.r_g[1:]**(l - 1)
            f_x = InterpolatedUnivariateSpline(
                self.r_g, a_g)(r_x) * r_x**(l - 1)

        f_x[:xcut] /= m_x[:xcut]
        f_k = fsbt(l, f_x, r_x, G_k)
        kcut = int(Gcut / G_k[1])
        f_k[kcut:] = 0.0
        ff_x = fsbt(l, f_k, G_k, r_x[:N // 2 + 1]) / pi * 2
        ff_x *= m_x[:N // 2 + 1]

        if l < 2:
            f_g = InterpolatedUnivariateSpline(
                r_x[:xcut + 1], ff_x[:xcut + 1])(self.r_g)
        else:
            ff_x[1:xcut + 1] /= r_x[1:xcut + 1]**(l - 1)
            f_g = InterpolatedUnivariateSpline(
                r_x[:xcut + 1], ff_x[:xcut + 1])(self.r_g) * self.r_g**(l - 1)
        f_g[self.ceil(rcut):] = 0.0

        return f_g

    def poisson(self, n_g, l=0):
        vr_g = self.zeros()
        nrdr_g = n_g * self.r_g * self.dr_g
        hartree(l, nrdr_g, self.r_g, vr_g)
        return vr_g

    def pseudize(self, a_g, gc, l=0, points=4):
        """Construct smooth continuation of a_g for g<gc.

        Returns (b_g, c_p[P-1]) such that b_g=a_g for g >= gc and::

                P-1      2(P-1-p)+l
            b = Sum c_p r
             g  p=0      g

        for g < gc+P.
        """
        assert isinstance(gc, numbers.Integral) and gc > 10, gc

        r_g = self.r_g
        i = np.arange(gc, gc + points)
        r_i = r_g[i]
        c_p = np.polyfit(r_i**2, a_g[i] / r_i**l, points - 1)
        b_g = a_g.copy()
        b_g[:gc] = np.polyval(c_p, r_g[:gc]**2) * r_g[:gc]**l
        return b_g, c_p[-1]

    def cut(self, a_g, rcut):
        gcut = self.floor(rcut)
        r0 = 0.7 * rcut
        x_g = np.clip((self.r_g - r0) / (rcut - r0), 0, 1)
        f_g = x_g**2 * (3 - 2 * x_g)
        shift = (4 * a_g[gcut] - a_g[gcut - 1]) / 3
        a_g -= f_g * shift
        a_g[gcut + 1:] = 0

    def pseudize_normalized(self, a_g, gc, l=0, points=3):
        """Construct normalized smooth continuation of a_g for g<gc.

        Same as pseudize() with also this constraint::

            /  _     2  /  _     2
            | dr b(r) = | dr a(r)
            /           /
        """

        b_g = self.pseudize(a_g, gc, l, points)[0]
        c_x = np.empty(points + 1)
        gc0 = gc // 2
        x0 = b_g[gc0]
        r_g = self.r_g
        i = [gc0] + list(range(gc, gc + points))
        r_i = r_g[i]
        norm = self.integrate(a_g**2)

        def f(x):
            b_g[gc0] = x[0]
            c_x[:] = np.polyfit(r_i**2, b_g[i] / r_i**l, points)
            b_g[:gc] = np.polyval(c_x, r_g[:gc]**2) * r_g[:gc]**l
            return self.integrate(b_g**2) - norm

        from scipy.optimize import fsolve
        fsolve(f, x0)
        return b_g, c_x[-1]

    def pseudize_smooth(self,
                        a_g: Array1D,
                        gc: int,
                        l: int = 0,
                        points: int = 3,
                        Gcut: float = 6.0) -> Tuple[Array1D, float]:
        """Minimize Fourier components above Gcut."""
        b_g, _ = self.pseudize(a_g, gc, l, points)
        c_x = np.empty(points + 1)
        gc0 = gc // 2
        x0 = b_g[gc0]
        i = [gc0] + list(range(gc, gc + points))
        r_g = self.r_g
        r_i = r_g[i]

        def f(x):
            b_g[gc0] = x[0]
            c_x[:] = np.polyfit(r_i**2, b_g[i] / r_i**l, points)
            b_g[:gc] = np.polyval(c_x, r_g[:gc]**2) * r_g[:gc]**l
            g_k, b_k = self.fft(b_g * r_g, l)
            kc = int(Gcut / g_k[1])
            f_k = g_k[kc:] * b_k[kc:]
            return f_k @ f_k

        from scipy.optimize import minimize
        minimize(f, x0)

        return b_g, c_x[-1]

    def jpseudize(self, a_g, gc, l=0, points=4):
        """Construct spherical Bessel function continuation of a_g for g<gc.

        Returns (b_g, b(0)/r^l) such that b_g=a_g for g >= gc and::

                P-2
            b = Sum c_p j (q r )
             g  p=0      l  p g

        for g < gc+P, where.
        """

        from scipy.optimize import brentq
        from scipy.special import sph_jn

        if a_g[gc] == 0:
            return self.zeros(), 0.0

        assert isinstance(gc, int) and gc > 10

        zeros_l = [[1, 2, 3, 4, 5, 6],
                   [1.430, 2.459, 3.471, 4.477, 5.482, 6.484],
                   [1.835, 2.895, 3.923, 4.938, 5.949, 6.956],
                   [2.224, 3.316, 4.360, 5.387, 6.405, 7.418]]

        # Logarithmic derivative:
        ld = np.dot([-1 / 60, 3 / 20, -3 / 4, 0, 3 / 4, -3 / 20, 1 / 60],
                    a_g[gc - 3:gc + 4]) / a_g[gc] / self.dr_g[gc]

        rc = self.r_g[gc]

        def f(q):
            j, dj = (y[-1] for y in sph_jn(l, q * rc))
            return dj * q - j * ld

        j_pg = self.empty(points - 1)
        q_p = np.empty(points - 1)

        zeros = zeros_l[l]
        if rc * ld > l:
            z1 = zeros[0]
            zeros = zeros[1:]
        else:
            z1 = 0
        x = 0.01
        for p, z2 in enumerate(zeros[:points - 1]):
            q = brentq(f, z1 * pi / rc + x, z2 * pi / rc - x)
            j_pg[p] = [sph_jn(l, q * r)[0][-1] for r in self.r_g]
            q_p[p] = q
            z1 = z2

        C_dg = [[0, 0, 0, 1, 0, 0, 0],
                [1 / 90, -3 / 20, 3 / 2, -49 / 18, 3 / 2, -3 / 20, 1 / 90],
                [1 / 8, -1, 13 / 8, 0, -13 / 8, 1, -1 / 8],
                [-1 / 6, 2, -13 / 2, 28 / 3, -13 / 2, 2, -1 / 6]][:points - 1]
        c_p = np.linalg.solve(np.dot(C_dg, j_pg[:, gc - 3:gc + 4].T),
                              np.dot(C_dg, a_g[gc - 3:gc + 4]))
        b_g = a_g.copy()
        b_g[:gc + 2] = np.dot(c_p, j_pg[:, :gc + 2])
        return b_g, np.dot(c_p, q_p**l) * 2**l * fac(l) / fac(2 * l + 1)

    def plot(self, a_g, n=0, rc=4.0, show=False):
        import matplotlib.pyplot as plt
        r_g = self.r_g[:len(a_g)]
        if n < 0:
            r_g = r_g[1:]
            a_g = a_g[1:]
        plt.plot(r_g, a_g * r_g**n)
        plt.axis(xmin=0, xmax=rc)
        if show:
            plt.show()

    def floor(self, r):
        return np.floor(self.r2g(r)).astype(int)

    def round(self, r):
        return np.around(self.r2g(r)).astype(int)

    def ceil(self, r):
        return np.ceil(self.r2g(r)).astype(int)

    def spline(self, a_g, rcut=None, l=0, points=None):
        """Create spline representation of a radial function f(r).

        The spline represents a rescaled version of the function, f(r) / r^l.
        """
        if points is None:
            points = self.default_spline_points

        if rcut is None:
            # Calculate rcut for the input function, i.e. a value rcut which
            # satisfies f(r) = 0 ∀ r ≥ rcut
            g = len(a_g) - 1
            while a_g[g] == 0.0:
                g -= 1
            rcut = self.r_g[g + 1]

        # Rescale f(r) -> f(r) / r^l
        N = len(a_g)
        b_g = divrl(a_g, l, self.r_g[:N])

        # Interpolate to a uniform radial grid (for the spline representation)
        r_i = np.linspace(0, rcut, points + 1)
        g_i = np.clip((self.r2g(r_i) + 0.5).astype(int), 1, N - 2)
        r1_i = self.r_g[g_i - 1]
        r2_i = self.r_g[g_i]
        r3_i = self.r_g[g_i + 1]
        x1_i = (r_i - r2_i) * (r_i - r3_i) / (r1_i - r2_i) / (r1_i - r3_i)
        x2_i = (r_i - r1_i) * (r_i - r3_i) / (r2_i - r1_i) / (r2_i - r3_i)
        x3_i = (r_i - r1_i) * (r_i - r2_i) / (r3_i - r1_i) / (r3_i - r2_i)
        b1_i = b_g[g_i - 1]
        b2_i = b_g[g_i]
        b3_i = b_g[g_i + 1]
        b_i = b1_i * x1_i + b2_i * x2_i + b3_i * x3_i

        return Spline(l, rcut, b_i)

    def get_cutoff(self, f_g):
        g = self.N - 1
        while f_g[g] == 0.0:
            g -= 1
        gcut = g + 1
        return gcut

    @abstractmethod
    def r2g(self, r):
        """Inverse continuous map from a real space distance (r)
           to a floating point grid index (g).

           Used by methods floor, round, and ceil, which manipulate this
           floating point to an integer accordingly.
        """


class EquidistantRadialGridDescriptor(RadialGridDescriptor):
    def __init__(self, h, N=1000, h0=0.0):
        """Equidistant radial grid descriptor.

        The radial grid is r(g) = h0 + g*h,  g = 0, 1, ..., N - 1."""

        RadialGridDescriptor.__init__(self,
                                      h * np.arange(N) + h0,
                                      h + np.zeros(N))

    def r2g(self, r):
        return int(np.ceil((r - self.r_g[0]) / (self.r_g[1] - self.r_g[0])))

    def xml(self, id='grid1'):
        assert self.r_g[0] == 0.0
        return ('<radial_grid eq="r=d*i" d="{0!r}" '
                'istart="0" iend="{1}" id="{2}"/>\n'
                .format(self.r_g[1], len(self.r_g) - 1, id))

    def spline(self, a_g, rcut=None, l=0, points=None):
        b_g = a_g.copy()
        if l > 0:
            b_g = divrl(b_g, l, self.r_g[:len(a_g)])
        return Spline(l, self.r_g[len(a_g) - 1], b_g)


class AERadialGridDescriptor(RadialGridDescriptor):
    def __init__(self, a, b, N=1000, default_spline_points=25):
        """Radial grid descriptor for all-electron calculation.

        The radial grid is::

                     a g
            r(g) = -------,  g = 0, 1, ..., N - 1
                   1 - b g
        """

        self.a = a
        self.b = b
        g = np.arange(N)
        r_g = self.a * g / (1 - self.b * g)
        dr_g = (self.b * r_g + self.a)**2 / self.a
        RadialGridDescriptor.__init__(self, r_g, dr_g, default_spline_points)

    def r2g(self, r):
        # return r / (r * self.b + self.a)
        # Hack to preserve backwards compatibility:
        ng = 1.0 / self.b
        beta = self.a / self.b
        return ng * r / (beta + r)

    def new(self, N):
        return AERadialGridDescriptor(self.a, self.b, N)

    def xml(self, id='grid1'):
        if abs(self.N - 1 / self.b) < 1e-5:
            return (('<radial_grid eq="r=a*i/(n-i)" a="%r" n="%d" ' +
                     'istart="0" iend="%d" id="%s"/>\n') %
                    (self.a * self.N, self.N, self.N - 1, id))
        return (('<radial_grid eq="r=a*i/(1-b*i)" a="%r" b="%r" n="%d" ' +
                 'istart="0" iend="%d" id="%s"/>\n') %
                (self.a, self.b, self.N, self.N - 1, id))

    def d2gdr2(self):
        return -2 * self.a * self.b / (self.b * self.r_g + self.a)**3


class AbinitRadialGridDescriptor(RadialGridDescriptor):
    def __init__(self, a, d, N=1000, default_spline_points=25):
        """Radial grid descriptor for Abinit calculations.

        The radial grid is::

                      dg
            r(g) = a(e   - 1),  g = 0, 1, ..., N - 1
        """

        self.a = a
        self.d = d
        g = np.arange(N)
        r_g = a * (np.exp(d * g) - 1)
        dr_g = (r_g + a) * d
        RadialGridDescriptor.__init__(self, r_g, dr_g, default_spline_points)

    def r2g(self, r):
        return np.log(r / self.a + 1) / self.d

    def d2gdr2(self):
        return -1 / (self.a**2 * self.d * (self.r_g / self.a + 1)**2)

    def new(self, N):
        return AbinitRadialGridDescriptor(self.a, self.d, N)
