# 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.

from bisect import insort

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 _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 _Interval:
    __slots__ = ['a', 'b', 'coeffs', 'vals', 'igral', 'err', 'level', 'depth',
                 'ndiv', 'c00', 'unreliable_err']

    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


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.attic = []         # Inactive intervals
        self.f = f
        self.igral_excess = 0
        self.err_excess = 0

    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."""
        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

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

        if ival.level == max_level:
            split = True
        else:
            points = self.refine(ival)
            split = ival.unreliable_err

            if (points[1] - points[0] < points[0] * min_sep
                or points[-1] - points[-2] < points[-2] * min_sep
                or ival.err < (abs(ival.igral)
                               * eps * tbls.V_cond_nums[ival.level])):
                # Remove the interval (while remembering the excess integral
                # and error), since it is either too narrow, or the estimated
                # relative error is already at the limit of numerical accuracy
                # and cannot be reduced further.
                self.err_excess += ival.err
                self.igral_excess += ival.igral
                self.attic.append(self.ivals.pop())
                return

        if split:
            # Replace current interval by its children.
            for new in self.split(self.ivals.pop()):
                insort(self.ivals, new)
        else:
            # The error estimate of the current interval has changed.
            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 vquad(f, a, b, tol):
    igrator = Vquad(f, a, b)
    return igrator.improve_until(tol)