import csv

import numpy as np
from numpy import diff

from scipy.interpolate import make_interp_spline
from scipy.signal import savgol_filter

import matplotlib.pyplot as plt

filenames = [
        'Laura/06_09-1.csv', # [0]
        'Laura/06_09-2.csv', # [1]
        'Elisabeth/06_10-1.csv', #[2]
        'Elisabeth/06_10-2.csv', #[3]
        # 'Elisabeth/06_12-1.csv',
        'Elisabeth/06_13-1.csv', #[4]
        'Elisabeth/06_13-2.csv', #[5]
        'Elisabeth/06_15-1.csv', #[6]
        'Elisabeth/06_15-2.csv', #[7]
        ]

comparison = {
        'Laura Shot #1 vs #2': (0, 1), #[0]
        'Elisabeth 06/10/22 Shot #2 - #1': (2, 3), #[1]
        'Elisabeth 06/13/22 Shot #2 - #1': (4, 5), #[2]
        'Elisabeth 06/15/22 Shot #2 - #1': (6, 7), #[3]
        'comparison': {
            'Elisabeth Shots 06/13/22 - 06/10/22': (1, 2), #[0]
            'Elisabeth Shots 06/15/22 - 06/10/22': (1, 3), #[1]
            'Elisabeth Shots 06/15/22 - 06/13/22': (2, 3), #[2]
            }
        }


data = {}

curr_figure = 0

plt.figure(curr_figure)
plt.subplot(2, 1, 1)
# plt.title('Shotweight over time')
# plt.xlabel('Time (s)')
plt.ylabel('Weight (g)')
plt.grid(visible=True, which='both', axis='both', linewidth=0.5)

plt.subplot(2, 1, 2)
# plt.title('Flow-rate over time')
plt.xlabel('Time (s)')
plt.ylabel('Flow-rate (g/s)')
plt.grid(visible=True, which='both', axis='both', linewidth=0.5)

plt.subplot(2, 1, 1)

for filename in filenames:
    data[filename] = {
            'shot_name': '',
            'rows': [],
            'time': [],
            'weight': []
            }

    with open(filename, newline='') as file:
        reader = csv.DictReader(file, delimiter=',')

        for row in reader:
            data[filename]['rows'].append(row)

            if row['information_type'] == 'meta':
                if row['metatype'] == 'Name':
                    data[filename]['shot_name'] = row['metadata']
            elif row['information_type'] == 'moment':
                try:
                    elapsed = float(row['elapsed'])
                    weight = float(row['current_total_shot_weight'])

                    data[filename]['time'].append(elapsed)
                    data[filename]['weight'].append(weight)
                except ValueError:
                    continue
                    #print("Not a float!")

    time = data[filename]['time']
    weight = data[filename]['weight']

    first_zero = list(x > 0 for x in weight).index(True) - 1
    zero_time = time[first_zero]

    weight = weight[first_zero:]
    time = time[first_zero:]
    time = list(map(lambda t: t - zero_time, time))

    data[filename]['time'] = time
    data[filename]['weight'] = weight

    weight_savgol = savgol_filter(weight, 30, 3)

    plt.plot(time, weight_savgol, label = data[filename]['shot_name'], linewidth=1)
    plt.scatter(time, weight, label = '', s=0.25)

    plt.subplot(2, 1, 2)

    dweight_dt = diff(weight)/diff(time)
    dweight_dt_savgol = savgol_filter(dweight_dt, 60, 3)

    plt.plot(time[:-1], dweight_dt_savgol, label = data[filename]['shot_name'], linewidth=1)
    plt.scatter(time[1:-1], dweight_dt[1:], label = '', s=0.25)

    plt.subplot(2, 1, 1)

    #print(filename, "\n", time, "\n", weight)

plt.legend(loc='best')

def comparison_plot(curr_figure, comparison_dict):
    curr_figure += 1
    plt.figure(curr_figure)

    for key in comparison_dict:
        if key != 'comparison':
            item = comparison_dict[key]
            i0 = item[0]
            i1 = item[1]

            fn0 = filenames[i0]
            fn1 = filenames[i1]

            plt.xlabel('Time (s)')
            plt.ylabel('Weight difference (g)')
            plt.grid(visible=True, which='both', axis='both', linewidth=0.5)

            t0 = np.array(data[fn0]['time'])
            w0 = np.array(data[fn0]['weight'])

            t1 = np.array(data[fn1]['time'])
            w1 = np.array(data[fn1]['weight'])

            time_linspace = np.linspace(max(np.array(t0).min(), np.array(t1).min()), min(np.array(t0).max(), np.array(t1).max()), max(len(t0), len(t1)))

            i0 = make_interp_spline(t0, w0, k=3)
            wi0 = i0(time_linspace)
            i1 = make_interp_spline(t1, w1, k=3)
            wi1 = i1(time_linspace)

            dw = wi1 - wi0

            dw_savgol = savgol_filter(dw, 30, 3)

            plt.plot(time_linspace, dw_savgol, label = key, linewidth=1)
            plt.scatter(time_linspace, dw, label = '', s=0.25)

            plt.legend(loc='best')
        else:
            curr_figure = comparison_plot(curr_figure, comparison_dict['comparison'])
    return curr_figure

curr_figure = comparison_plot(curr_figure, comparison)


plt.show()