import pandas as pd

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

shots = [
        {
            'filename': 'Elisabeth/16_06_2022_1.xlsx',
            'title': 'Elisabeth 16.06.2022 #1',
            'cutoff': 38,
            'data': {}
        }, # [0]
        ]

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 shot in shots:
    df_raw = pd.read_excel(shot['filename'], converters={
        1: lambda x: float(x) if not str(x).endswith('.10') else float(x) + 0.85
        }, sheet_name=0)
    df_calc = pd.read_excel(shot['filename'], converters={
        1: lambda x: float(x) if not str(x).endswith('.10') else float(x) + 0.85
        }, sheet_name=1)

    time_col_raw = df_raw.keys()[1]
    profile_col_raw = df_raw.keys()[5]

    time_col_calc = df_calc.keys()[1]
    profile_col_calc = df_calc.keys()[2]

    df_raw_co = df_raw.loc[df_raw[time_col_raw] < shot['cutoff']]
    df_calc_co = df_calc.loc[df_calc[time_col_calc] < shot['cutoff']]

    shot['data']['profile'] = df_raw_co
    shot['data']['flowrate'] = df_calc_co

    time = df_raw_co[time_col_raw].tolist()
    weight = df_raw_co[profile_col_raw].tolist()
    print("time: ", time)
    print("weight: ", time)
    plt.plot(time, weight, label = shot['title'], linewidth=1)

    plt.subplot(2, 1, 2)

    flow_time = df_calc_co[time_col_calc].tolist()
    flowrate = df_calc_co[profile_col_calc].tolist()
    print("flow_time: ", flow_time)
    print("flowrate: ", flowrate)
    plt.plot(flow_time, flowrate, label = shot['title'], linewidth=1)

    plt.subplot(2, 1, 1)

plt.legend(loc='best')

plt.show()