# Copyright 2017, 2018 Christoph Groth (CEA).
#
# This file is part of Vquad.  It is subject to the license terms in the file
# LICENSE.rst found in the top-level directory of this distribution.

import bisect

import numpy as np
from scipy.linalg import norm

from . import tables as tbls

eps = np.spacing(1)

# If the relative difference between two consecutive approximations is
# lower than this value, the error estimate is considered reliable.
# See section 6.2 of Pedro Gonnet's thesis.
hint = 0.1

# Smallest acceptable relative difference of points in a rule.  This was chosen
# such that no artifacts are apparent in plots of (i, log(a_i)), where a_i is
# the sequence of estimates of the integral value of an interval and all its
# ancestors..
min_sep = 16 * eps

min_level = 1
max_level = 4

ndiv_max = 20

_sqrt_one_half = np.sqrt(0.5)


def _eval_legendre(c, x):
    """Evaluate _orthonormal_ Legendre polynomial.

    This uses the three-term recurrence relation from page 63 of Perdo
    Gonnet's thesis.
    """
    if len(c) <= 1:
        c0 = c[0]
        c1 = 0
    else:
        n = len(c)
        c0 = c[-2]              # = c[k + 0]
        c1 = c[-1]              # = c[k + 1]
        for k in range(len(c) - 2, 0, -1):
            a = (2*k + 3) / (k + 1)**2
            tmp = c0
            c0 = c[k - 1] - c1 * np.sqrt(a * k**2 / (2*k - 1))
            c1 = tmp + c1 * x * np.sqrt(a * (2*k + 1))

    return np.sqrt(1/2) * c0 + np.sqrt(3/2) * c1 * x


def _calc_coeffs(vals, level):
    nans = np.flatnonzero(~np.isfinite(vals))
    if nans.size:
        # Replace vals by a copy and zero-out non-finite elements.
        vals = vals.copy()
        vals[nans] = 0
        # Prepare things for the loop further down.
        b = tbls.newton_coeffs[level].copy()
        m = len(b) - 2              # = len(tbls.nodes[level]) - 1
    coeffs = tbls.inv_Vs[level] @ vals

    # This is a variant of Algorithm 7 from the thesis of Pedro Gonnet where no
    # linear system has to be solved explicitly.  Instead, Algorithm 5 is used.
    for i in nans:
        b[m + 1] /= tbls.alpha[m]
        x = tbls.nodes[level][i]
        b[m] = (b[m] + x * b[m + 1]) / tbls.alpha[m - 1]
        for j in range(m - 1, 0, -1):
            b[j] = ((b[j] + x * b[j + 1] - tbls.gamma[j + 1] * b[j + 2])
                    / tbls.alpha[j - 1])
        b = b[1:]

        coeffs[:m] -= coeffs[m] / b[m] * b[:m]
        coeffs[m] = 0
        m -= 1

    return coeffs


class DivergentIntegralError(ValueError):
    def __init__(self, msg, igral, err):
        self.igral = igral
        self.err = err
        super().__init__(msg)


class _Terminator:
    __slots__ = ['prev', 'next']


class _Interval:
    __slots__ = ['a', 'b', 'coeffs', 'vals', 'igral', 'err', 'level', 'depth',
                 'ndiv', 'c00', 'unreliable_err', 'prev', 'next']

    def __init__(self, a, b, level, depth):
        self.a = a
        self.b = b
        self.level = level
        self.depth = depth

    def points(self):
        a = self.a
        b = self.b
        return (a + b) / 2 + (b - a) * tbls.nodes[self.level] / 2

    def interpolate(self, vals, coeffs_old=None):
        self.vals = vals
        self.coeffs = coeffs = _calc_coeffs(self.vals, self.level)
        if self.level == min_level:
            self.c00 = coeffs[0]
        if coeffs_old is None:
            coeffs_diff = norm(coeffs)
        else:
            coeffs_diff = np.zeros(max(len(coeffs_old), len(coeffs)))
            coeffs_diff[:len(coeffs_old)] = coeffs_old
            coeffs_diff[:len(coeffs)] -= coeffs
            coeffs_diff = norm(coeffs_diff)
        w = self.b - self.a
        self.igral = w * coeffs[0] * _sqrt_one_half
        self.err = w * coeffs_diff
        self.unreliable_err = coeffs_diff > hint * norm(coeffs)

    def __lt__(self, other):
        return self.err < other.err

    def __call__(self, x):
        a = self.a
        b = self.b
        x = (2 * x - (a + b)) / (b - a)
        return _eval_legendre(self.coeffs, x)

class Vquad:
    """Evaluate an integral using adaptive quadrature.

    The algorithm uses Clenshaw-Curtis quadrature rules of increasing
    degree in each interval.  The error estimate is
    sqrt(integrate((f0(x) - f1(x))**2)), where f0 and f1 are two
    successive interpolations of the integrand.  To fall below the
    desired total error, intervals are worked on ranked by their own
    absolute error: either the degree of the rule is increased or the
    interval is split if either the function does not appear to be
    smooth or a rule of maximum degree has been reached.

    Reference: "Increasing the Reliability of Adaptive Quadrature
        Using Explicit Interpolants", P. Gonnet, ACM Transactions on
        Mathematical Software, 37 (3), art. no. 26, 2008.
    """

    def __init__(self, f, a, b, level=max_level - 1):
        ival = _Interval(a, b, level, 1)
        vals = f(ival.points())
        ival.interpolate(vals)
        ival.c00 = 0.0          # Will go away.
        ival.ndiv = 0

        self.ivals = [ival]     # Active intervals
        self.f = f
        self.igral_excess = 0
        self.err_excess = 0

        # Initialize linked list.
        ival.prev = self.begin = _Terminator()
        self.begin.next = ival
        ival.next = self.end = _Terminator()
        self.end.prev = ival

    def split(self, ival):
        m = (ival.a + ival.b) / 2
        f_center = ival.vals[(len(ival.vals) - 1) // 2]

        depth = ival.depth + 1
        children = [_Interval(ival.a, m, min_level, depth),
                    _Interval(m, ival.b, min_level, depth)]
        points = np.concatenate([child.points()[1:-1] for child in children])
        valss = np.empty((2, tbls.sizes[min_level]))
        valss[:, 0] = ival.vals[0], f_center
        valss[:, -1] = f_center, ival.vals[-1]
        valss[:, 1:-1] = self.f(points).reshape((2, -1))

        for child, vals, T in zip(children, valss, tbls.Ts):
            child.interpolate(vals, T[:, :ival.coeffs.shape[0]] @ ival.coeffs)
            child.ndiv = (ival.ndiv
                         + (ival.c00 and child.c00 / ival.c00 > 2))
            if child.ndiv > ndiv_max and 2*child.ndiv > child.depth:
                msg = ('Possibly divergent integral in the interval '
                       '[{}, {}]! (h={})')
                raise DivergentIntegralError(
                    msg.format(child.a, child.b, child.b - child.a),
                    child.igral * np.inf, None)
        return children

    def refine(self, ival):
        """Increase degree of interval.

        Returns True if the refined interval is OK, and False if it is
        borderline and should not be refined/split further.  This
        happens when neigboring points can be only barely resolved by
        floating point numbers, or when the estimated relative error is
        already at the limit of numerical accuracy and cannot be reduced
        further.
        """
        ival.level += 1
        points = ival.points()
        vals = np.empty(points.shape)
        vals[0::2] = ival.vals
        vals[1::2] = self.f(points[1::2])
        ival.interpolate(vals, ival.coeffs)

        return (points[1] - points[0] > points[0] * min_sep
                and points[-1] - points[-2] > points[-2] * min_sep
                and ival.err > (abs(ival.igral)
                                * eps * tbls.V_cond_nums[ival.level]))

    def improve(self):
        ival = self.ivals[-1]

        if ival.level == max_level:
            split = True
        else:
            if not self.refine(ival):
                # Remove the interval but remember the excess integral and
                # error.
                self.err_excess += ival.err
                self.igral_excess += ival.igral
                self.ivals.pop()
                return
            split = ival.unreliable_err

        if split:
            # Replace current interval by its children.
            self.ivals.pop()
            child0, child1 = self.split(ival)
            bisect.insort(self.ivals, child0)
            bisect.insort(self.ivals, child1)

            # Maintain linked list.
            ival.prev.next = child0
            ival.next.prev = child1
            child0.prev = ival.prev
            child0.next = child1
            child1.prev = child0
            child1.next = ival.next
        else:
            # The error estimate of the current interval has changed.
            bisect.insort(self.ivals, self.ivals.pop())

    def totals(self):
        igral = self.igral_excess
        err = self.err_excess
        for ival in self.ivals:
            igral += ival.igral
            err += ival.err
        return igral, err

    def improve_until(self, tol):
        while True:
            self.improve()
            igral, err = self.totals()

            if (err == 0
                or err < abs(igral) * tol
                or err - self.err_excess < abs(igral) * tol < self.err_excess
                or not self.ivals):
                return igral, err

    def __call__(self, xs):
        xs = np.asarray(xs)
        shape = xs.shape
        if xs.size == 0:
            return np.empty(shape)
        xs = xs.flatten()

        # Sort xs, but remember inverse permutation.
        perm = np.argsort(xs)
        xs = xs[perm]
        inv_perm = np.empty(len(perm), int)
        inv_perm[perm] = np.arange(len(perm))

        # Evaluate points interval by interval.
        results = []
        ival = self.begin.next
        end = self.end
        if xs[0] < ival.a or xs[-1] > end.prev.b:
            raise ValueError("Point lies outside of integration interval.")
        i = 0
        while ival is not end:
            j = bisect.bisect(xs, ival.b, i)
            if j != i:
                results.append(ival(xs[i:j]))
                i = j
            ival = ival.next

        return np.concatenate(results)[inv_perm].reshape(shape)


def vquad(f, a, b, tol):
    igrator = Vquad(f, a, b)
    return igrator.improve_until(tol)