"""
Created on 28. aug. 2015

@author: pab
"""

from __future__ import annotations

import warnings
from typing import Any

import numpy as np
from numpy.typing import ArrayLike
from scipy import linalg
from scipy.ndimage import convolve1d

from numdifftools._typing import Array, ArrayOrScalar, ExtrapolatedSequence

FINFO: np.finfo[Any] = np.finfo(float)
EPS: float = FINFO.eps
_tiny_name = "tiny" if np.__version__ < "1.22" else "smallest_normal"
_TINY: float = getattr(FINFO, _tiny_name)
_HUGE: float = FINFO.max


def _assert(cond: bool, msg: str) -> None:
    if not cond:
        raise ValueError(msg)


def convolve(sequence: ArrayLike, rule: ArrayLike, **kwds: Any) -> Array:
    """Wrapper around scipy.ndimage.convolve1d that allows complex input."""
    dtype = np.result_type(float, np.ravel(sequence)[0])
    seq = np.asarray(sequence, dtype=dtype)
    if np.iscomplexobj(seq):
        return convolve1d(seq.real, rule, **kwds) + 1j * convolve1d(seq.imag, rule, **kwds)
    return convolve1d(seq, rule, **kwds)


class Dea:
    """
    Extrapolate a slowly convergent sequence using repeated Shanks transformations.

    Notes
    -----
    DEA attempts to extrapolate nonlinearly by Shanks transformations to a better
    estimate of the sequence's limiting value, thus improving the rate of convergence.
    The epsilon algorithm of P. Wynn, see [1]_, is used to perform the
    non-linear Shanks transformations. The routine is a translation of the
    DQELG function found in the QUADPACK fortran library, see [2]_ and [3]_.

    List of major variables:

    LIMEXP:  scalar integer
        The maximum number of elements the epsilon table data can contain.
        The epsilon table is stored in the first (LIMEXP+2) entries of EPSTAB.
    EPSTAB: real vector or size (LIMEXP+2+3)
        The first LIMEXP+2 elements contains the two lower diagonals of the triangular
        epsilon table. The elements are numbered starting at the right-hand corner of the
       triangle.
    E0,E1,E2,E3:  real scalars
        The 4 elements on which the computation of a new element in the epsilon table is based.
    NRES:  scalar integer
        Number of extrapolation results actually generated by the epsilon algorithm in prior
        calls to the routine.
    NEWELM: scalar integer
        Number of elements to be computed in the new diagonal of the epsilon table.
        The condensed epsilon table is computed. Only those elements needed for the
        computation of the next diagonal are preserved.
    RES: real scalar
        New element in the new diagonal of the epsilon table.
    ERROR: real scalar
        An estimate of the absolute error of RES. The routine decides whether RESULT=RES or
        RESULT=SVALUE by comparing ERROR with abserr from the previous call.
    RES3LA: real vector of size 3
        Contains at most the last 3 results.

    Reference
    ---------
    ..  [1] Wynn, P. (1956)
            "On a Device for Computing the em(Sn) Transformation",
            Mathematical Tables and Other Aids to Computation, 10, 91-96.
    ..  [2] R. Piessens, E. De Doncker-Kapenga and C. W. Uberhuber (1983),
            "QUADPACK: a subroutine package for automatic integration",
            Springer, ISBN: 3-540-12553-1, 1983.
    ..  [3] http://www.netlib.org/quadpack/
    ..  [4] https://mathworld.wolfram.com/WynnsEpsilonMethod.html
    """

    def __init__(self, limexp: int = 50) -> None:
        self.limexp = limexp
        self._n = 0
        self._nres = 0

    @property
    def limexp(self) -> int:
        """Maximum number of elements the epsilon table data."""
        return self._limexp

    @limexp.setter
    def limexp(self, limexp: int) -> None:
        n = 2 * (limexp // 2) + 1
        _assert(n >= 3, "LIMEXP IS LESS THAN 3")
        self.epstab = np.zeros(n + 5)
        self._limexp = n

    def _dea(self, epstab: Array, n: int) -> tuple[ArrayOrScalar, float, int]:
        res3la = epstab[-3:]
        nres = self._nres

        abserr = _HUGE
        result = epstab[n]
        limexp = self.limexp
        epstab[n + 2] = epstab[n]
        newelm = n // 2
        epstab[n] = _HUGE
        old_n = n
        k_1 = n
        all_converged = False
        for i in range(newelm):
            res = epstab[k_1 + 2]
            e_0 = epstab[k_1 - 2]
            e_1 = epstab[k_1 - 1]
            e_2 = res
            delta2 = e_2 - e_1
            delta3 = e_1 - e_0
            err2 = abs(delta2)
            err3 = abs(delta3)
            e1abs = abs(e_1)
            tol2 = max(abs(e_2), e1abs) * EPS
            tol3 = max(e1abs, abs(e_0)) * EPS
            all_converged = not (err2 > tol2 or err3 > tol3)
            if all_converged:
                # if e_0, e_1 and e_2 are equal to within machine accuracy, convergence is assumed.
                result = res
                abserr = err2 + err3

                # ***jump out of do-loop
                # go to 100
                break
            e_3 = epstab[k_1]
            epstab[k_1] = e_1
            delta1 = e_1 - e_3
            err1 = abs(delta1)
            tol1 = max(e1abs, abs(e_3)) * EPS

            # if two elements are very close to each other, omit
            #  a part of the table by adjusting the value of n

            any_converged = err1 <= tol1 or err2 <= tol2 or err3 <= tol3
            if not any_converged:  # go to 20
                sss = 1.0 / delta1 + 1.0 / delta2 - 1.0 / delta3
                epsinf = abs(sss * e_1)

                # If irregular behaviour in the table is detected, omit a
                # part of the table adjusting the value of n.
                any_converged = epsinf <= 1e-4

            if any_converged:
                n = 2 * i
                # ***jump out of do-loop
                # go to 50
                break

            #  compute a new element and eventually adjust the value of result.
            res = e_1 + 1.0 / sss
            epstab[k_1] = res
            k_1 = k_1 - 2
            error = err2 + abs(res - e_2) + err3
            if not error > abserr:
                abserr = error
                result = res
            # 40 continue

        # 50
        if not all_converged:
            #      shift the table.
            if n == limexp - 1:
                n = limexp - 2  # 2*(limexp//2) - 1
            self._shift_table(epstab, n, newelm, old_n)
            if nres > 1:
                abserr = np.abs(result - res3la[:nres]).sum()

            self._update_res3la(res3la, result, nres)
        # 100
        abserr = max(abserr, 5.0 * EPS * abs(result))
        self._nres += 1
        return result, float(abserr), n

    @staticmethod
    def _shift_table(epstab: Array, n: int, newelm: int, old_n: int) -> Array:
        i_0 = old_n % 2  # 1 if ((old_n // 2) * 2 == old_n - 1) else 0
        i_n = 2 * newelm + 2
        epstab[i_0:i_n:2] = epstab[i_0 + 2 : i_n + 2 : 2]

        if old_n != n:
            i_n = old_n - n
            epstab[: n + 1] = epstab[i_n : i_n + n + 1]
        return epstab

    @staticmethod
    def _update_res3la(res3la: Array, result: ArrayOrScalar, nres: int) -> None:
        if nres > 2:
            res3la[:2] = res3la[1:]
            res3la[2] = result
        else:
            res3la[nres] = result

    def __call__(self, s_value: ArrayOrScalar) -> tuple[ArrayOrScalar, float]:
        epstab = self.epstab

        result = s_value
        n = self._n

        epstab[n] = s_value
        if n == 0:
            abserr = abs(result)
        elif n == 1:
            abserr = 6.0 * np.abs(result - epstab[0])
        else:
            result, abserr, n = self._dea(epstab, n)
        n += 1
        self._n = n

        abserr = np.asarray(abserr).item()
        return result, float(abserr)


class EpsAlg:
    """
    Extrapolate a slowly convergent sequence using Shanks transformation.

    Notes
    -----
    The iterated Shanks transformation is computed using the Wynn
    epsilon algorithm (see equation 4.3-10a to 4.3-10c given on page 25 in [1]_).


    References
    ----------
    ..  [1] E. J. Weniger (1989)
            "Nonlinear sequence transformations for the acceleration of
            convergence and the summation of divergent series"
            Computer Physics Reports Vol. 10, 189 - 371
            http://arxiv.org/abs/math/0306302v1
    ..  [2] https://mathworld.wolfram.com/WynnsEpsilonMethod.html
    """

    def __init__(self) -> None:
        self.epstab: list[ArrayOrScalar] = []

    def __call__(self, s_n: ArrayOrScalar) -> ArrayOrScalar:
        epstab = self.epstab
        n = len(epstab)
        epstab.append(s_n)
        if n == 0:
            estlim = s_n
        else:
            aux1: ArrayOrScalar = 0.0
            aux2: ArrayOrScalar = 0.0

            for i in range(n, 0, -1):
                aux1 = aux2
                aux2 = epstab[i - 1]
                delta = epstab[i] - aux2
                if np.abs(delta) <= 1.0e-60:
                    epstab[i - 1] = 1.0e60
                else:
                    epstab[i - 1] = aux1 + 1.0 / delta
            estlim = epstab[n % 2]

        return estlim


def linfun(i: int) -> Array:
    return np.linspace(0, np.pi / 2.0, 2**i + 1)


def richardson_demo() -> None:
    """
    >>> from numdifftools.extrapolation import richardson_demo
    >>> richardson_demo()
    NO. PANELS      TRAP. APPROX          APPROX W/R            abserr
        1           0.78539816            0.78539816            0.21460184
        2           0.94805945            1.11072073            0.11072073
        4           0.98711580            0.99798929            0.00201071
        8           0.99678517            0.99988201            0.00011799
       16           0.99919668            0.99999274            0.00000726
       32           0.99979919            0.99999955            0.00000045
       64           0.99994980            0.99999997            0.00000003
      128           0.99998745            1.00000000            0.00000000
      256           0.99999686            1.00000000            0.00000000
      512           0.99999922            1.00000000            0.00000000
    """

    n = 10
    e_i = []
    h = []

    print("NO. PANELS      TRAP. APPROX          APPROX W/R            abserr")
    txt = "{0:5d} {1:20.8f}  {2:20.8f}  {3:20.8f}"
    for k in range(n):
        x = linfun(k)
        val = np.trapezoid(np.sin(x), x)
        h.append(x[1])
        e_i.append(val)
        vale, _err0, _step = Richardson(step=1, order=1)(np.array(e_i), np.array(h))

        err = np.abs(1.0 - vale)
        print(txt.format(len(x) - 1, val, vale[-1], err[-1]))


def epsalg_demo() -> None:
    """
    >>> from numdifftools.extrapolation import epsalg_demo
    >>> epsalg_demo()
    NO. PANELS      TRAP. APPROX          APPROX W/EA           abserr
        1           0.78539816            0.78539816            0.21460184
        2           0.94805945            0.94805945            0.05194055
        4           0.98711580            0.99945672            0.00054328
        8           0.99678517            0.99996674            0.00003326
       16           0.99919668            0.99999988            0.00000012
       32           0.99979919            1.00000000            0.00000000
       64           0.99994980            1.00000000            0.00000000
      128           0.99998745            1.00000000            0.00000000
      256           0.99999686            1.00000000            0.00000000
      512           0.99999922            1.00000000            0.00000000
    """

    dea = EpsAlg()
    print("NO. PANELS      TRAP. APPROX          APPROX W/EA           abserr")
    txt = "{0:5d} {1:20.8f}  {2:20.8f}  {3:20.8f}"
    for k in range(10):
        x = linfun(k)
        val = np.trapezoid(np.sin(x), x)
        vale = dea(val)
        err = np.abs(1.0 - vale)
        print(txt.format(len(x) - 1, val, vale, err))


def dea_demo() -> None:
    """
    >>> from numdifftools.extrapolation import dea_demo
    >>> dea_demo()
    NO. PANELS      TRAP. APPROX          APPROX W/EA           abserr
        1           0.78539816            0.78539816            0.78539816
        2           0.94805945            0.94805945            0.97596771
        4           0.98711580            0.99945672            0.21405856
        8           0.99678517            0.99996674            0.05190729
       16           0.99919668            0.99999988            0.00057629
       32           0.99979919            1.00000000            0.00057665
       64           0.99994980            1.00000000            0.00003338
      128           0.99998745            1.00000000            0.00000012
      256           0.99999686            1.00000000            0.00000000
      512           0.99999922            1.00000000            0.00000000
     1024           0.99999980            1.00000000            0.00000000
     2048           0.99999995            1.00000000            0.00000000
    """

    dea = Dea(limexp=6)
    print("NO. PANELS      TRAP. APPROX          APPROX W/EA           abserr")
    txt = "{0:5d} {1:20.8f}  {2:20.8f}  {3:20.8f}"
    vals = []
    num_panels = []
    for k in range(12):
        x = linfun(k)
        val = np.trapezoid(np.sin(x), x)
        vals.append(val)
        num_panels.append(len(x) - 1)
    for k, val in zip(num_panels, vals, strict=False):
        vale, err = dea(val)
        print(txt.format(k, val, vale, err))


def max_abs(a: ArrayLike, b: ArrayLike) -> Array:
    """Returns element-wise maximum of absulute value of array elements"""
    return np.maximum(np.abs(a), np.abs(b))


def dea3(v_0: ArrayLike, v_1: ArrayLike, v_2: ArrayLike, symmetric: bool = False) -> tuple[Array, Array]:
    """
    Extrapolates a slowly convergent sequence using Shanks transformations.

    Parameters
    ----------
    v_0, v_1, v_2 : array-like
        Three consecutive values of a convergent sequence to extrapolate.
    symmetric : bool, optional
        If True, returns a sequence with symmetric error estimates.
        Defaults to False.

    Returns
    -------
    result : array
        The extrapolated value(s).
    abserr : array
        The estimated absolute error(s).

    Notes
    -----
    DEA3 uses nonlinear Shanks transformations based on three values to
    extrapolate to a better estimate of the sequence's limit. It is a vectorized
    translation of the DQELG function from the QUADPACK Fortran library.
    for LIMEXP=3, see [2]_ and [3]_.

    Examples
    --------
    # integrate sin(x) from 0 to pi/2

    >>> import numpy as np
    >>> import numdifftools as nd
    >>> Ei= np.zeros(3)
    >>> linfun = lambda i : np.linspace(0, np.pi/2., 2**(i+5)+1)
    >>> for k in range(3):
    ...    x = linfun(k)
    ...    Ei[k] = np.trapezoid(np.sin(x),x)
    >>> [En, err] = nd.dea3(Ei[0], Ei[1], Ei[2])
    >>> truErr = np.abs(En-1.)
    >>> bool(np.all(truErr < err))
    True
    >>> np.allclose(En, 1)
    True
    >>> bool(np.all(np.abs(Ei-1)<1e-3))
    True

    See also
    --------
    Dea

    References
    ----------
    ..  [1] Wynn, P. (1956)
            "On a Device for Computing the em(Sn) Transformation",
            Mathematical Tables and Other Aids to Computation, 10, 91-96.
    ..  [2] R. Piessens, E. De Doncker-Kapenga and C. W. Uberhuber (1983),
            "QUADPACK: a subroutine package for automatic integration",
            Springer, ISBN: 3-540-12553-1, 1983.
    ..  [3] http://www.netlib.org/quadpack/
    ..  [4] https://mathworld.wolfram.com/WynnsEpsilonMethod.html
    """
    e_0, e_1, e_2 = np.atleast_1d(v_0, v_1, v_2)

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)  # ignore division by zero
        delta2, delta1 = e_2 - e_1, e_1 - e_0

        # Calculate absolute errors and tolerances
        err2, err1 = np.abs(delta2), np.abs(delta1)
        tol2, tol1 = max_abs(e_2, e_1) * EPS, max_abs(e_1, e_0) * EPS

        # Avoid division by zero
        delta1[err1 < _TINY] = _TINY
        delta2[err2 < _TINY] = _TINY
        # The core Shanks transformation logic: 1 / (1/delta2 - 1/delta1)
        s_inv_diff = 1.0 / delta2 - 1.0 / delta1

        # Check for convergence based on different criteria
        relative_error_small = abs(s_inv_diff * e_1) <= 1.0e-4
        converged = (err1 <= tol1) | (err2 <= tol2) | relative_error_small

        # Apply the transformation only where not converged
        result = np.where(converged, e_2, e_1 + 1.0 / s_inv_diff)

        # Estimate the error, using a conservative approach where convergence is met
        abserr_term = np.where(converged, tol2 * 10, np.abs(result - e_2))
        abserr = err1 + err2 + abserr_term

    if symmetric and len(result) > 1:
        return result[:-1], abserr[1:]
    return result, abserr


class Richardson:
    """
    Extrapolates a sequence with Richardsons method

    Parameters
    ----------
    step_ratio: real scalar
        Ratio between sequential steps, h, generated.
    step: scalar integer
        Defines the step between exponents in the error polynomial,
        i.e., step = k_1 - k_0 = k_2 - k_1 = ... = k_{i+1} - k_i
    order: scalar integer
        Leading order of truncation error.
    num_terms: scalar integer
        Number of terms used in the polynomial fit.

    Notes
    -----
    Suppose f(h) is an approximation of L (exact value) that depends on a positive
    step size h described with a sequence of the form

    L = f(h) + a0 * h^k_0 + a1 * h^k_1+ a2 * h^k_2 + ...

    where the ai are unknown constants and the k_i are known constants such that h^k_i > h^(k_i+1).

    If we evaluate the right hand side for different stepsizes h
    we can fit a polynomial to that sequence of approximations.
    This is exactly what this class does.
    Here k_0 is the leading order step size behavior of truncation error as L = f(h)+O(h^k_0)
    (f(h) -> L as h -> 0, but f(0) != L) and k_i = order + step * i .

    Examples
    --------
    >>> import numpy as np
    >>> import numdifftools as nd
    >>> n = 3
    >>> Ei = np.zeros((n,1))
    >>> h = np.zeros((n,1))
    >>> linfun = lambda i : np.linspace(0, np.pi/2., 2**(i+5)+1)
    >>> for k in range(n):
    ...    x = linfun(k)
    ...    h[k] = x[1]
    ...    Ei[k] = np.trapezoid(np.sin(x),x)
    >>> En, err, step = nd.Richardson(step=1, order=1)(Ei, h)
    >>> truErr = np.abs(En-1.)
    >>> bool(np.all(truErr < err))
    True
    >>> bool(np.all(np.abs(Ei-1)<1e-3))
    True
    >>> np.allclose(En, 1)
    True
    """

    step_ratio: float
    step: int | float
    order: int
    num_terms: int

    def __init__(
        self,
        step_ratio: float = 2.0,
        step: int | float = 1,
        order: int = 1,
        num_terms: int = 2,
    ) -> None:
        self.step_ratio = step_ratio
        self.step = step
        self.order = order
        self.num_terms = num_terms

    @staticmethod
    def _r_matrix(step_ratio: float, step: int | float, num_terms: int, order: int) -> Array:
        i, j = np.ogrid[0 : num_terms + 1, 0:num_terms]
        dtype = np.result_type(step_ratio, step, float)
        r_mat = np.ones((num_terms + 1, num_terms + 1), dtype=dtype)
        r_mat[:, 1:] = (1.0 / step_ratio) ** (i * (step * j + order))
        return r_mat

    def rule(
        self,
        sequence_length: int | None = None,
    ) -> Array:
        """Returns extrapolation rule."""
        if sequence_length is None:
            sequence_length = self.num_terms + 1
        num_terms = min(self.num_terms, sequence_length - 1)
        if num_terms > 0:
            r_mat = self._r_matrix(self.step_ratio, self.step, num_terms, self.order)
            return linalg.pinv(r_mat)[0]
        return np.ones((1,))

    @staticmethod
    def _estimate_error(new_sequence: Array, old_sequence: Array, steps: Array, rule: Array) -> Array:
        m = new_sequence.shape[0]
        m_old = old_sequence.shape[0]
        cov1 = np.sum(np.abs(rule) ** 2)  # 1 spare dof
        fact = np.maximum(12.7062047361747 * np.sqrt(cov1), EPS * 10.0)
        if m_old < 2:
            return (np.abs(new_sequence) * EPS + steps) * fact
        if m < 2:
            delta = np.diff(old_sequence, axis=0)
            tol = max_abs(old_sequence[:-1], old_sequence[1:]) * fact
            err = np.abs(delta)
            converged = err <= tol
            abserr = err[-m:] + np.where(
                converged[-m:], tol[-m:] * 10, abs(new_sequence - old_sequence[-m:]) * fact
            )
            return abserr
        #         if m_old>2:
        #             res, abserr = dea3(old_sequence[:-2], old_sequence[1:-1],
        #                               old_sequence[2:] )
        #             return abserr[-m:] * fact
        err = np.abs(np.diff(new_sequence, axis=0)) * fact
        tol = max_abs(new_sequence[1:], new_sequence[:-1]) * EPS * fact
        converged = err <= tol
        abserr = err + np.where(converged, tol * 10, abs(new_sequence[:-1] - old_sequence[-m + 1 :]) * fact)
        return abserr

    def extrapolate(self, sequence: Array, steps: Array) -> ExtrapolatedSequence:
        """Extrapolate sequence"""
        return self.__call__(sequence, steps)

    def __call__(
        self,
        sequence: Array,
        steps: Array,
    ) -> ExtrapolatedSequence:
        num_steps = sequence.shape[0]
        rule = self.rule(num_steps)
        n_r = rule.size - 1
        m = num_steps - n_r
        k = min(num_steps, m + 1)
        new_sequence = convolve(sequence, rule[::-1], axis=0, origin=n_r // 2)
        abserr = self._estimate_error(new_sequence[:k], sequence, steps, rule)
        return ExtrapolatedSequence(new_sequence[:m], abserr[:m], steps[:m])


if __name__ == "__main__":
    from numdifftools.testing import test_docstrings

    test_docstrings(__file__)
