"""
Problem 2
"""

import numpy as np
from scipy import integrate as quad
from scipy.optimize import minimize_scalar
from matplotlib import pyplot as plt


def f(x):
    return np.exp(-np.cos(x))


def g(x):
    return f(x) * (np.heaviside(x+1, 0) - np.heaviside(x-1, 0))


def plot_fg():
    plt.figure()
    xn = np.linspace(-np.pi, np.pi, 500)
    plt.plot(xn, f(xn))
    plt.plot(xn, g(xn))


def test_blackbox():
    I1 = quad.romberg(f, -np.pi, np.pi)
    I2 = quad.romberg(f, -1, 1)
    I3 = quad.romberg(g, -np.pi, np.pi)
    print(f'I1 = {I1}\nI2 = {I2}\nI3 = {I3}')


def trap(func, lower, upper, num_points):
    xn, dx = np.linspace(lower, upper, num_points, retstep=True)
    yn = func(xn)
    return dx * (yn[0]/2 + yn[-1]/2 + np.sum(yn[1:-1])), dx


def conv_rate(func, lower, upper, N_base):
    A1, _ = trap(func, lower, upper, N_base)
    A2, _ = trap(func, lower, upper, 2 * N_base)
    A3, _ = trap(func, lower, upper, 3 * N_base)
    return minimize_scalar(lambda k: np.abs(A2 + (A2-A1)/(2**k - 1) - A3 - (A3-A1)/(3**k - 1))).x


def test_trap():
    params = [(f, -np.pi, np.pi, 7.9549265210128452745132196653294, r'$I_1$', 5),
              (f, -1, 1, 0.87070265620795901020832433774759, r'$I_2$', 100),
              (g, -np.pi, np.pi, 0.87070265620795901020832433774759, r'$I_3$', 100)]
    N_vals = 2**np.arange(11, dtype='i')
    plt.figure()
    ax1 = plt.subplot(111,
                      title='Error vs step size',
                      xscale='log', yscale='log',
                      xlabel='Step size', ylabel='Absolute error'
                      )
    for func, a, b, ans, name, N_conv in params:
        # Calculate simple error
        h_vals = np.zeros(N_vals.shape[0])
        error = np.zeros(N_vals.shape[0])
        for i, N in enumerate(N_vals):
            value, h_vals[i] = trap(func, a, b, N)
            error[i] = np.abs(value - ans)
        ax1.plot(h_vals, error, label=name)
        # Find convergence using Richardson extrapolation
        conv = conv_rate(func, a, b, N_conv)
        print(f'Convergence of {name} : O(h^{conv})')
    plt.legend()


if __name__ == '__main__':
    print('COMMIT!!!')
    plot_fg()
    test_blackbox()
    test_trap()
    plt.show()

