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


infile  =  'random-poly.xlsx'
outfile  = 'lsq-polynomial-out.csv'

fontsize = 16


argv = sys.argv
narg = len(argv)
if narg >= 2:
    infile = argv[1]


def lsq(x, y, iPrint = 0):
    n = len(x)

# numpy ndarrayでは、ベクトルも 2x1の 行列としないと行列計算ができない
    Si  = np.empty([2, 1])
    Sij = np.empty([2, 2])

    x_sum  = 0.0
    x2_sum = 0.0
    y_sum  = 0.0
    xy_sum = 0.0
    for i in range(0, n):
        x_sum  += x[i]
        x2_sum += x[i]**2
        y_sum  += y[i]
        xy_sum += x[i] * y[i]

    Sij[0, 0], Sij[1, 0] = n, x_sum
    Sij[0, 1], Sij[1, 1] = x_sum, x2_sum
    Si[0], Si[1] = y_sum, xy_sum

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

# 2x2行列と2x1行列の積なので、返り値の ci は2x1行列
    ci = np.linalg.inv(Sij) @ Si
# ciの転置を取ってリスト変数に変換
    ci = ci.transpose().tolist()

# ciの転置の第０成分が係数ベクトル
    return ci[0]


def main():
    print("Least-squares method for linear function")
    print(f"norder=1")
    print(f"infile={infile}")
    print(f"outfile={outfile}")

    print("")
    print(f"Read [{infile}]")
    df = pd.read_excel(infile, engine = 'openpyxl')
    labels = df.columns.to_list()
    x = df[labels[0]]
    y = df[labels[1]]
    ndata = len(x)

    print("")
    print(f"Execute linear least-squares method")
    ci = lsq(x, y, iPrint = 1)

    print("LSQ function")
    print(f"f(x) = {ci[0]} + {ci[1]} * x")

    print("")
    ycal = []
    f = open(outfile, 'w')
    fout = csv.writer(f, lineterminator='\n')
    fout.writerow(['x', 'y', 'y(LSQ)'])
    for i in range(0, ndata):
        yl = ci[0] + ci[1] * x[i]
        ycal.append(yl)
        fout.writerow([x[i], y[i], yl])

#================================================================
# Plot
#================================================================
    fig, axes = plt.subplots(1, 1, figsize = (8, 6))
    axes.plot(x, y,    label = 'input', linestyle = '', marker = 'o')
    axes.plot(x, ycal, label = 'fit',   linestyle = '-')
    axes.tick_params(labelsize = fontsize)
    axes.set_xlabel('$x$', fontsize = fontsize)
    axes.set_ylabel('$y$', fontsize = fontsize)
    axes.legend(fontsize = fontsize)

    plt.tight_layout()
    plt.pause(0.1)

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


if __name__ == "__main__":
    main()