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


"""
scipy.optimize.leastsqを使って多項式最小二乗の結果をグラフにプロット
フィッティング範囲を起動時引数で変える
"""


#=============================
# 大域変数の定義
#=============================
# CSVファイル
infile = 'data.csv'

# フィッティングパラメータ初期値。線形最小二乗の場合は適当
ai0 = [0, 0, 0]

fitrange = [-1e10, 1e10]

# グラフのフォントサイズ
fontsize = 24

# 起動時引数でフィッティング範囲を変更
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("")
print("fitting range: ", fitrange)

#=============================
# 最小化する関数の定義
#=============================
def ycal(ai, x):
    return ai[0] + ai[1] * x + ai[2] * x * x

def residual(ai, x, y):
    res = []
    for i in range(len(x)):
        res.append(y[i] - ycal(ai, x[i]))
    return res
    

#=============================
# 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:
            xi = float(row[0])
            if fitrange[0] <= xi <= fitrange[1]:
                x.append(xi)
                y.append(float(row[1]))
        i = i + 1

print("")
print("CSV data:")
print("  header:", header)
print("  x:", x)
print("  y:", y)


#=============================
# scipy.optimize()による最小化
#=============================
print("")
print("polynomial fit by scipy.optimize() start:")
# leastsqの戻り値は、最適化したパラメータのリストと、最適化の結果
ai, cov_x, inf, mesg, ier = optimize.leastsq(residual, ai0, args= (x, y), full_output = True)
print("  lsq result: ai=", ai)
res = sqrt(sum(inf['fvec']*inf['fvec']) / len(x))
print("              residual=", res)
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
    yi = ycal(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.title(infile, fontsize = fontsize)
plt.xlabel(header[0], fontsize = fontsize)
plt.ylabel(header[1], fontsize = fontsize)
plt.legend(fontsize = fontsize)
plt.tick_params(labelsize = fontsize)
plt.tight_layout()

plt.show()
