# 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 numpy as np
from numpy.testing import assert_allclose
from pytest import raises

from .. import core

def f0(x):
    return x * np.sin(1/x) * np.sqrt(abs(1 - x))


def f7(x):
    return x**-0.5


def f24(x):
    return np.floor(np.exp(x))


def f21(x):
    y = 0
    for i in range(1, 4):
        y += 1 / np.cosh(20**i * (x - 2 * i / 10))
    return y


def f63(x, alpha, beta):
    return abs(x - beta) ** alpha


def F63(x, alpha, beta):
    return (x - beta) * abs(x - beta) ** alpha / (alpha + 1)


def fdiv(x):
    return abs(x - 0.987654321) ** -1.1


def f_one_with_nan(x):
    x = np.asarray(x)
    result = np.ones(x.shape)
    result[x == 0] = np.inf
    return result


def test_coeffs(level=3):
    vals = np.abs(core.tbls.nodes[level])
    vals[1::2] = np.nan
    c_downdated = core._calc_coeffs(vals, level)

    level -= 1
    vals = np.abs(core.tbls.nodes[level])
    c = core._calc_coeffs(vals, level)
    assert_allclose(c_downdated[:len(c)], c, rtol=0, atol=1e-12)

    vals_from_c = core._eval_legendre(c, core.tbls.nodes[level])
    assert_allclose(vals_from_c, vals, rtol=0, atol=1e-15)


def test_integration():
    old_settings = np.seterr(all='ignore')

    igral, err = core.vquad(f0, 0, 3, 1e-5)
    assert_allclose(igral, 1.98194117954329, 1e-15)
    assert_allclose(err, 1.9563545589988155e-05, 1e-10)

    igral, err = core.vquad(f7, 0, 1, 1e-6)
    assert_allclose(igral, 1.9999998579359648, 1e-15)
    assert_allclose(err, 1.8561437334964041e-06, 1e-10)

    igral, err = core.vquad(f24, 0, 3, 1e-3)
    assert_allclose(igral, 17.664696186312934, 1e-15)
    assert_allclose(err, 0.017602618074957457, 1e-10)

    igral, err = core.vquad(f21, 0, 1, 1e-3)
    assert_allclose(igral, 0.16310022131213361, 1e-15)
    assert_allclose(err, 0.00011848806384952786, 1e-10)

    igral, err = core.vquad(f_one_with_nan, -1, 1, 1e-12)
    assert_allclose(igral, 2, 1e-15)
    assert_allclose(err, 2.4237853822937613e-15, 1e-7)

    try:
        igral, err = core.vquad(fdiv, 0, 1, 1e-6)
    except core.DivergentIntegralError as e:
        assert e.igral == np.inf
        assert e.err is None

    np.seterr(**old_settings)


def test_interpolation():
    vquad = core.Vquad(f24, 0, 3)
    vquad.improve_until(1e-6)

    rng = np.random.RandomState(123)
    x = np.linspace(0, 3, 100)
    rng.shuffle(x)

    for x in [x, 1.23, [[2, 0], [1, 2]], [], [[]]]:
        assert_allclose(vquad(x), f24(x))

    for x in [-1e-100, 3.0001, 1e100, [1, 2, -1e50]]:
        raises(ValueError, vquad, x)


def test_analytic(n=200):
    def f(x):
        return f63(x, alpha, beta)

    def F(x):
        return F63(x, alpha, beta)

    old_settings = np.seterr(all='ignore')

    np.random.seed(123)
    params = np.empty((n, 2))
    params[:, 0] = np.linspace(-0.5, -1.5, n)
    params[:, 1] = np.random.random_sample(n)

    false_negatives = 0
    false_positives = 0

    for alpha, beta in params:
        try:
            igral, err = core.vquad(f, 0, 1, 1e-3)
        except core.DivergentIntegralError:
            assert alpha < -0.8
            false_negatives += alpha > -1
        else:
            if alpha <= -1:
                false_positives += 1
            else:
                igral_exact = F(1) - F(0)
                assert alpha < -0.7 or abs(igral - igral_exact) < err

    assert false_negatives < 0.05 * n
    assert false_positives < 0.05 * n

    np.seterr(**old_settings)