import sys
import numpy as np
from numpy import sin, cos, tan, pi
from pprint import pprint
import csv
import openpyxl
import pandas as pd
import matplotlib.pyplot as plt


norders = [2, 4, 6, 8, 12]

x0 = -2.0
x1 =  2.0

xcal0 = -2.5
xcal1 =  2.5
ndata = 101

fontsize = 16


def func(x):
    return 1.0 / (1.0 + x * x)

def mlsq(x, y, m, iPrint = 0):
    n = len(x)
    Si  = np.empty([m+1, 1])
    Sij = np.empty([m+1, m+1])

    for l in range(m+1):
        v = 0.0
        for i in range(n):
            v += y[i] * pow(x[i], l)
        Si[l, 0] = v

    for j in range(m+1):
        for l in range(j, m+1):
            v = 0.0
            for i in range(n):
                v += pow(x[i], j+l)
            Sij[j, l] = Sij[l, j] = v

    if iPrint == 1:
        print("Vector and Matrix:")
        print("Si=")
        pprint(Si)
        print("Sij=")
        pprint(Sij)
        print("")

    ci = np.linalg.inv(Sij) @ Si
    ci = ci.transpose().tolist()
    return ci[0]


def main():
    print("Least-squares method for polynomial order {}".format(norders))

    xcalstep = (xcal1 - xcal0) / (ndata - 1)
    xcal = [xcal0 + i * xcalstep for i in range(ndata)]
    ycals = []
    for norder in norders:
        print("")
        print(f"norder={norder}")

        print(f"Make func data")
        xstep = (x1 - x0) / norder
        x = [x0 + i * xstep for i in range(norder+1)]
        y = []
        for i in range(norder + 1):
            y.append(func(x[i]))

        ci = mlsq(x, y, norder, iPrint = 0)

        print("  LSQ function")
        print("    f(x) = {:6.3g}".format(ci[0]), end = '')
        for i in range(1, norder+1):
            print(" + {:6.3g} * x^{}".format(ci[i], i), end = '')
        print("")

        ycal = []
        for i in range(ndata):
            _x = xcal[i]
            yl = ci[0]
            for k in range(1, norder+1):
                yl += ci[k] * pow(_x, k)

            ycal.append(yl)

        ycals.append(ycal)

#================================================================
# Plot
#================================================================
    fig, axes = plt.subplots(1, 1, figsize = (8, 6))
    axes.plot(x, y,       label = 'input', linestyle = '-', linewidth = 1.0, marker = 'o', markersize = 4.0)
    for i in range(len(norders)):
        axes.plot(xcal, ycals[i], label = f'order {norders[i]}',   linestyle = '-', linewidth = 0.5)
    axes.tick_params(labelsize = fontsize)
    axes.set_xlabel('$x$', fontsize = fontsize)
    axes.set_ylabel('$y$', fontsize = fontsize)
    axes.set_ylim([0.0, 1.2])
    axes.legend(fontsize = 8)

    plt.tight_layout()
    plt.pause(0.1)

    print("")
    print("Press ENTER to terminate")
    input()


if __name__ == "__main__":
    main()