"""
Problem 1:
"""
import time
from scipy.special import factorial
import numpy as np
from matplotlib import pyplot as plt


class SineExpand:
    things = 'terms'

    def __init__(self, num_terms):
        t0 = time.perf_counter()
        n = np.arange(num_terms)
        self.coeff = (-1)**n / factorial(2*n+1)
        self.power = 2*n + 1
        t1 = time.perf_counter()
        self.overhead = t1 - t0

    def evaluate(self, x):
        """x must be a 1D numpy array"""
        t0 = time.perf_counter()
        coeff = np.tile(self.coeff, (x.shape[0], 1)).T
        power = np.tile(self.power, (x.shape[0], 1)).T
        t1 = time.perf_counter()
        result = np.sum(coeff[:] * x**power[:], axis=0)
        t2 = time.perf_counter()
        rmserror = np.sqrt(np.mean((np.sin(x)-result)**2))
        overhead = self.overhead
        evaltime = t2 - t0
        return rmserror, evaltime/x.shape[0], overhead


class SineInterp:
    things = 'points'

    def __init__(self, num_terms):
        t0 = time.perf_counter()
        self.x = np.linspace(0, 2*np.pi, num_terms)
        self.y = np.sin(self.x)
        t1 = time.perf_counter()
        self.overhead = t1 - t0

    def evaluate(self, x):
        """x must be a 1D numpy array"""
        t0 = time.perf_counter()
        result = np.interp(x, self.x, self.y)
        t1 = time.perf_counter()
        rmserror = np.sqrt(np.mean((np.sin(x)-result)**2))
        evaltime = t1 - t0
        return rmserror, evaltime/x.shape[0], self.overhead


def test_sine(sine_class, N_vals):
    x_vals = np.random.default_rng().random(10**6) * 2 * np.pi
    rmserror = np.zeros(N_vals.shape[0])
    evaltime = np.zeros(N_vals.shape[0])
    overhead = np.zeros(N_vals.shape[0])
    for i, N in enumerate(N_vals):
        sine = sine_class(N)
        rmserror[i], evaltime[i], overhead[i] = sine.evaluate(x_vals)
    plt.figure()
    plt.subplot(121, title='Accuracy', xlabel=f'Number of {sine_class.things}', ylabel='RMS Error')
    plt.yscale('log')
    plt.plot(N_vals, rmserror)
    plt.subplot(122, title='Speed', xlabel=f'Number of {sine_class.things}', ylabel='Time (sec)')
    plt.plot(N_vals, evaltime, label='Evaluation')
    plt.plot(N_vals, overhead, label='Overhead')
    plt.yscale('log')
    plt.legend()
    return rmserror, evaltime, overhead


if __name__ == '__main__':
    print('COMMIT!!!')
    err1, et1, ot1 = test_sine(SineExpand, np.arange(1, 11))
    err2, et2, ot2 = test_sine(SineInterp, np.geomspace(3, 500, 100, dtype='i'))
    plt.figure()
    plt.subplot(121, title='Error vs evaluation time', ylabel='Evaluation time (s)', xlabel='Absolute error',
                xscale='log', yscale='log')
    plt.plot(err1, et1, label='Taylor expansion')
    plt.plot(err2, et2, label='Linear interpolation')
    plt.legend()
    plt.subplot(122, title='Error vs overhead time', ylabel='Overhead time (s)', xlabel='Absolute error',
                xscale='log', yscale='log')
    plt.plot(err1, ot1, label='Taylor expansion')
    plt.plot(err2, ot2, label='Linear interpolation')
    plt.legend()
    plt.show()
