import sys
import os
from tkinter import *
from tkinter import ttk
from tkinter import filedialog, messagebox
from pprint import pprint
import csv
from math import sqrt
import numpy as np
from scipy import optimize
from matplotlib import pyplot as plt


"""
GUIでファイル選択をして最小二乗
"""


#=============================
# 大域変数の定義
#=============================
# フィッティングパラメータ初期値。線形最小二乗の場合は適当
ai0 = [0, 0, 0]
xrange = [0, 10]
# グラフのフォントサイズ
fontsize = 24

window_size = "500x200"
ini_dir = os.path.abspath(os.path.dirname(__file__))
file_type = [('CSV file', '*.csv')]


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

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

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


def path_button_click(pathvar, x0var, x1var, obj):
    selpath = filedialog.askopenfilename(filetypes = file_type, initialdir = ini_dir)
    pathvar.set(selpath)
    header, x, y = read_csv(selpath)
    x0var.set(min(x))
    x1var.set(max(x))
    obj['path']   = selpath
    obj['header'] = header
    obj['x'] = x
    obj['y'] = y
    obj['savepath'].set(os.path.splitext(selpath)[0] + '-fit.csv')

#=============================
# 最小化する関数の定義
#=============================
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


#=============================
# scipy.optimize()による最小化
#=============================
def lsq_button_click(x, y, x0, x1, obj):
    print("")
    print("polynomial fit by scipy.optimize() start:")

    xf = []
    yf = []
    for i in range(len(x)):
        if x0 <= x[i] <= x1:
            xf.append(x[i])
            yf.append(y[i])

# leastsqの戻り値は、最適化したパラメータのリストと、最適化の結果
    ai, cov_x, inf, mesg, ier = optimize.leastsq(residual, ai0, 
            args= (xf, yf), 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]))


#=============================
# グラフの表示
#=============================
#表示データの作成
    ncal = 100
    xmin = min(x)
    xmax = max(x)
    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.clf()
    plt.plot(x,  y,  label = 'raw data', marker = 'o', linestyle = 'None')
    plt.plot(xc, yc, label = 'fitted', linestyle = 'dashed')
    plt.title(obj['path'], fontsize = fontsize)
    plt.xlabel(obj['header'][1], fontsize = fontsize)
    plt.ylabel(obj['header'][0], fontsize = fontsize)
    plt.legend(fontsize = fontsize)
    plt.tick_params(labelsize = fontsize)
    plt.tight_layout()

    plt.pause(0.001)
#    plt.show()


def main():
    obj = {}

    root = Tk()
    root.title('Second-order polynomial lsq')
#    root.resizable(False, False)
    root.geometry(window_size)
    root.minsize(200, 200)

    path     = StringVar()
    savepath = StringVar()
    x0 = DoubleVar(value = xrange[0])
    x1 = DoubleVar(value = xrange[1])
    obj['savepath'] = savepath

# Menu
    menu_bar = Menu(root)
    menu_file = Menu(menu_bar, tearoff = 0)
    menu_file.add_command(label='Open', accelerator='Ctrl+O',
                    command = lambda: path_button_click(path, x0, x1, obj))
    menu_file.add_command(label='exit', accelerator='Alt+E',
                    command = lambda: exit())
    menu_bar.add_cascade(label = 'File', menu = menu_file)
    root.config(menu = menu_bar)
    root.grid()

# Root frame
    root_frame = ttk.Frame(root, padding=10)
    root_frame.pack(side = 'top')

# Path frame
    path_frame = ttk.Frame(root_frame)
    path_label = ttk.Label(path_frame, text = 'Path:', padding = (5,2))
    path_label.pack(side = 'left')
    path_entry = ttk.Entry(
        path_frame,
        textvariable = path,
        width = 50
        )
    path_entry.pack(side = 'left', expand = True)

    path_button = ttk.Button(path_frame, text = 'path', 
                    command = lambda: path_button_click(path, x0, x1, obj))
    path_button.pack(side = 'left')
    path_frame.pack(side = 'top', anchor = 'w')

# Save path frame
    savepath_frame = ttk.Frame(root_frame)
    savepath_label = ttk.Label(savepath_frame, text = 'Save path:', padding = (5,2))
    savepath_label.pack(side = 'left')
    savepath_entry = ttk.Entry(
        savepath_frame,
        textvariable = savepath,
        width = 50
        )
    savepath_entry.pack(side = 'left', expand = True)

    savepath_button = ttk.Button(savepath_frame, text = 'save', 
                    command = lambda: savepath_button_click(path, obj))
    savepath_button.pack(side = 'left')
    savepath_frame.pack(side = 'top', anchor = 'w')

# Range frame
    range_frame = ttk.Frame(root_frame)
    range_label = ttk.Label(range_frame, text = 'x range:', padding = (5,2))
    range_label.grid(row = 1, column = 0, sticky = 'w')
    x0_entry = ttk.Entry(
        range_frame,
        textvariable = x0,
        width = 10)
    x0_entry.grid(row = 1, column = 1)

    range_label2 = ttk.Label(range_frame, text = '-', padding = (5,2))
    range_label2.grid(row = 1, column = 2, sticky = E)
    x1_entry = ttk.Entry(
        range_frame,
        textvariable = x1,
        width = 10)
    x1_entry.grid(row = 1, column = 3)
    range_frame.pack(side = 'top', anchor = 'w')

# Button frame
    button_frame = ttk.Frame(root_frame)
    lsq_button = ttk.Button(button_frame, text = 'lsq',
        command = lambda: lsq_button_click(obj['x'], obj['y'], x0.get(), x1.get(), obj))
    lsq_button.pack(side = 'left')
    exit_button = ttk.Button(button_frame, text = 'exit', command = lambda: exit())
    exit_button.pack(side = 'left')
    button_frame.pack(side = 'top', anchor = 'w')
    
    root.mainloop()


if __name__ == '__main__':
    main()
    