#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt
import time

def poly_fit(x, y, order):
    """ 
    Fits the coefficients of a polynomial function to 
    a set of two-dimensional data points using the least-squares method.

    Note: uses the function 'numpy.linalg.lstsq' which returns a tuple, where the first item is the solution vector.

    Parameters
    ----------
    x : list
        x-coordinates of the data points
    y : list
        y-coordinates of the data points
    order: int
        Order of the polynomial to fit the data, see eq. (6) in 'main.ipynb' for the polynomial form:
        https://sgit.iue.tuwien.ac.at/360049/homework8/src/branch/main/main.ipynb#user-content-Aufgabe-3:-Ausgleichungsrechnung-mit-Polynomen-(1-Punkt)

    Returns
    -------
    list
        Coefficients of the polynomial
    """
    A = np.array([[xi**n for n in range(order+1)] for xi in x])
    b = np.array(y)
    start_time = time.perf_counter()
    # coeff,res,rank,s = np.linalg.lstsq(A, b, rcond=None)
    coeff = np.linalg.solve(A.T@A, A.T@b)
    end_time = time.perf_counter()
    execution_time_ms = float((end_time - start_time) * 1000)
    print(f"N={len(x)} runtime: {execution_time_ms}ms")    
    return coeff


def plot_plot(x, y, func, filename):

    plt.figure()

    plt.plot(x, y, marker="o", linestyle="", label="Data")

    x_samples = np.linspace(min(x),max(x),100)
    yp = func(x_samples)
    
    plt.plot(x_samples, yp, label="Fitting")

    plt.xlabel("x")
    plt.ylabel("y")
    plt.legend()
    plt.savefig(filename)


if __name__ == "__main__":


    
    N = 30000


    # w/o noise, , fit n=3
    x = np.linspace(0, 5, N)
    y = 2*(1 - np.exp(-x)) #  y = 2*(1 - exp(-x))
    n = 3
    coeff = poly_fit(x, y, n)
    print(coeff)
    
    # note: plotting disabled 
    # func = lambda x : sum([coeff[i]*x**i for i in range(0, n+1)])
    # plot_plot(x, y, func, "taskA_wo_noise_n3.png")

    # exponential w/ noise, fit n=3
    mean = 0.0
    sigma = 0.2
    y = y + np.random.normal(mean, sigma, len(y))
    n = 3
    coeff = poly_fit(x, y, n)
    print(coeff)

    # note: plotting disabled 
    # func = lambda x : sum([coeff[i]*x**i for i in range(0, n+1)])
    # plot_plot(x, y, func, "taskA_w_noise_n3.png")

    # exponential w/ noise, fit n=4
    n = 4
    coeff = poly_fit(x, y, n)
    print(coeff)

    # note: plotting disabled 
    # func = lambda x : sum([coeff[i]*x**i for i in range(0, n+1)])
    # plot_plot(x, y, func, "taskA_w_noise_n4.png")