# numpy.polyfitを使って多項式最小二乗の結果をグラフにプロットしてみる
# http://ailaby.com/least_square/
# https://qiita.com/Morio/items/d75159bac916174e7654
# https://rikei-fufu.com/2019/06/28/post-1357-python-plot/

import sys
import csv
from pprint import pprint
from math import sqrt
import numpy as np
from matplotlib import pyplot as plt

#=============================
# 大域変数の定義
#=============================
infile = 'data.csv'
fitrange = [0, 2.0]

argv = sys.argv
print("argv=", argv)

if len(argv) >= 2:
    fitrange[0] = float(argv[1])
if len(argv) >= 3:
    fitrange[1] = float(argv[2])

print("fitting range: ", fitrange)


#=============================
# csvファイルの読み込み
#=============================
i = 0
x = []
y = []
with open(infile, "r") as f:
    reader = csv.reader(f)

    for row in reader:
        if i == 0:
            header = row
        else:
            x.append(float(row[0]))
            y.append(float(row[1]))
        i = i + 1

print("header:", header)

print("x:", x)
print("y:", y)


#=======================================
# numpy.polyfit()による多項式最小二乗
#=======================================
print("polynomial fit start:")
ret = np.polyfit(x, y, deg = 2, full = True)
print("ret=", ret)
ai = ret[0]
residual = sqrt(ret[1][0] / len(x))
print("  lsq result: ai=", ai)
print("              residual=", residual)
#print("  y = {} + {} * x + {} * x^2".format(ai[0], ai[1], ai[2]))

xmin = min(x)
xmax = max(x)
ncal = 100
xstep = (xmax - xmin) / (ncal - 1)
xc = []
yc = []
for i in range(ncal):
    xi = xmin + i * xstep
    if fitrange[0] <= xi <= fitrange[1]:
        yi = np.poly1d(ai)(xi)
        xc.append(xi)
        yc.append(yi)

plt.plot(x, y, label='raw data', marker = 'o', linestyle = 'None')
plt.plot(xc, yc, label='fitted', linestyle = 'dashed')
plt.xlabel(header[0])
plt.ylabel(header[1])
plt.legend()

plt.show()
