#!/usr/bin/env python3

# Composite Numerical Integration:
# Trapezoid and Simpson Rules
#---------------------------------

import numpy as np
import scipy.integrate as integrate

def trapz(yk, dx):
    """ Trapezoid Integration """
    longsum = 2*np.sum(yk[1:-1])

    return dx/2 * (longsum + yk[0] + yk[-1])

def test_trapz():
        def test_func(x):
            return np.exp(x**2 - x**3)

        for _ in range(10):
            N = 1e1
            a = np.random.rand(1)
            b = np.random.rand(1)

            xk = np.linspace(a, b, N)

            yk = test_func(xk)
            h = (b-a)/N

            assert np.isclose( trapz(yk, h) , integrate.trapz(yk, dx=h), rtol=1/N)




    
def simps(yk, dx):
    """ Integration using Simpson's 1/3 Rule """
    even_part = 2*np.sum(yk[2:-2:2])
    odd_part = 4*np.sum(yk[1:-1:2])

    return dx/3 * ( even_part + odd_part + yk[0] + yk[-1] )






def main_2a():
    N = 1e3
    a = 0
    b = 2

    xk, h = np.linspace(a, b, N, retstep=True)

    func = lambda x: np.exp(- x**2 - 1 )

    yk = func(xk)

    print("My Integrators")
    print("Trapz:", trapz(yk,h))
    print("Simps:", simps(yk,h))
    print()
    print("Scipy Integrators")
    print("Trapz:", integrate.trapz(yk,dx=h))
    print("Simps:", integrate.simps(yk,dx=h))


def main_2b():
    import pytest
    pytest.main(__file__)


def main_2c():
    from matplotlib import pyplot

    N = 10
    funcs = [ lambda x: x, lambda x: x**2, lambda x: x**(1/2) ]
    exact = [ 1/2, 1/3, 3/2 ]
    diffs = np.zeros((np.size(funcs), N, 2))

    Ns = np.logspace(1,4, N)
    for i, f in enumerate(funcs):
        for j, n in enumerate(Ns):
            xk,h = np.linspace(0, 1, n, retstep=True)
            yk = f(xk)

            print(i,j)
            diffs[i][j][0] = exact[j] - trapz(yk,h)
            diffs[i][j][1] = exact[j] - simps(yk,h)


    pyplot.subplots()
    pyplot.title("Difference between the exact and numeric solution")
    pyplot.xlabel("N")
    pyplot.ylabel("difference")
    pyplot.plot(Ns, diffs[0][:][0], label="Trapz f(x)=x")
    
    pyplot.show()





if __name__ == "__main__":
    main_2a()
    #main_2b()
    #main_2c()