# https://www.cnblogs.com/vachester/p/7202793.html
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn import datasets, linear_model
import argparse


def read_data(file_name, start_row):
    data = pd.read_csv(file_name, header=None, sep='\t')
    # data = pd.read_excel(file_name)

    row_bin1 = []
    wave_bin1 = []
    for i in range(start_row, start_row + 300):
        row_bin1.append([float(i)])
    for i in range(data.shape[1]):
        wave_bin1.append(data.iloc[0][i])

    row_bin2 = []
    wave_bin2 = []
    for i in np.arange(start_row / 2, start_row / 2 + 150):
        row_bin2.append([float(i)])
    for i in range(0, len(wave_bin1), 2):
        # print(i)
        # print(wave[i:i + 2])
        wave_bin2.append((wave_bin1[i] + wave_bin1[i+1])/2)

    return row_bin1, wave_bin1, row_bin2, wave_bin2


def plot(x, y, regre):
    plt.scatter(x, y, color='blue')
    plt.plot(x, regre.predict(x), color='red', linewidth=4)
    # plt.xticks(())
    # plt.yticks(())
    plt.show()


def linearRegression(X_parameters, Y_parameters):#
    regr = linear_model.LinearRegression()
    regr.fit(X_parameters, Y_parameters)

    # 绘图
    # plot(X_parameters, Y_parameters, regr)

    return regr


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("csv_path", help="Path of csv file which contains wavelength.")
    parser.add_argument("start_row", help="Start row of coning 410 sensor.")
    args = parser.parse_args()

    row_bin1, wave_bin1, row_bin2, wave_bin2 = read_data(args.csv_path, int(args.start_row))
    regr_bin1 = linearRegression(row_bin1, wave_bin1)
    regr_bin2 = linearRegression(row_bin2, wave_bin2)

    # bin1 calculate
    yCalculated_bin1 = []
    a_bin1 = regr_bin1.coef_[0]
    b_bin1 = regr_bin1.intercept_
    for i in range(len(row_bin1)):
        yCalculated_bin1.append(a_bin1 * row_bin1[i][0] + b_bin1)

    # bin2 calculate
    yCalculated_bin2 = []
    a_bin2 = regr_bin2.coef_[0]
    b_bin2 = regr_bin2.intercept_
    for i in range(len(row_bin2)):
        yCalculated_bin2.append(a_bin2 * row_bin2[i][0] + b_bin2)

    #
    df_bin1 = pd.DataFrame(list(zip(wave_bin1, yCalculated_bin1)),
                           columns=['original', 'Calculated_bin1'])
    df_bin2 = pd.DataFrame(list(zip(yCalculated_bin2)), columns=['Calculated_bin2'])

    path_tmp = args.csv_path.split(".")
    path_new = path_tmp[0] + "_result.xlsx"
    writer = pd.ExcelWriter(path_new)
    df_bin1.to_excel(excel_writer=writer, sheet_name="bin1",
                     columns=["original", "Calculated_bin1"], index=False)
    df_bin2.to_excel(excel_writer=writer, sheet_name="bin2",
                     columns=["Calculated_bin2"], index=False)
    writer.save()
    writer.close()

    print("bin1")
    print("Intercept value ", regr_bin1.intercept_)
    print("coefficient", regr_bin1.coef_)
    print("bin2")
    print("Intercept value ", regr_bin2.intercept_)
    print("coefficient", regr_bin2.coef_)

    print("completed!!")