# file to evaluate a measurement
import os.path, glob
from os.path import join
import json

import pandas as pd

# from bec_analysis.settings import GROUP_STORAGE_PATH, LOCAL_GROUP_STORAGE
import numpy as np
import matplotlib.pyplot as pl
# from bec_analysis.evaluation.shapiro_steps.shapiro_steps_batch import *
# from bec_experiment.bec_style import *
# from bec_analysis.calibration.barrier_height_calibration import *
# from bec_analysis.calibration.chemical_potential import *
#
from helper import *
from style import *


# path to data
path_data_repo = "../Fig_1"


# def get_meas_data(meas):
#
#     path_exp = join(path_data_repo, meas["path"])
#     # Read Experiment Data
#     data = pd.read_csv(path_exp)
#     # data = data[data["discard"] != True]
#     # data = data[data["sum_fit_offsetcomp"] >= 50000]
#     # # remove reference
#     # data = data[data[SEGMENT_MULTIPLIER] == 1]
#     # data = data[data[Z_REF] > -1]
#
#
#
#     data["chem_pot"] = mu_N(data["N"] * (1 + data['delta_z'])) - mu_N(
#         data["N"] * (1 - data['delta_z']))
#
#     print(len(data["chem_pot"]))
#
#     f = data["f"].to_numpy()[0]
#
#     I_m = data["I_m"].to_numpy()[0] # I_m / I_c
#
#     I_c = current(1, v_c=1)
#
#     # if store:
#     #     # save to data repo
#     #     df = pd.DataFrame(data={'vel': data[V]*scaling,
#     #                             'delta_z': data[DELTA_Z],
#     #                             'N': data["sum_fit_offsetcomp"],
#     #                             'I_m': I_m/I_c,
#     #                             'f': f})
#     #     df = df.set_index('vel')
#     #
#     #     df.to_csv(os.path.join(path_data_repo, name_csv))
#
#     data = data[["delta_z", "vel", "chem_pot"]]
#     grouped = data.groupby(by="vel").agg(['mean', 'sem'])
#
#     vel = grouped.index.to_numpy()
#     z = grouped['delta_z', 'mean']
#     z_err = grouped['delta_z', 'sem']
#
#     mu = grouped['chem_pot', 'mean']
#     mu_err = grouped['chem_pot', 'sem']
#
#     I = current(vel, v_c=0.42)
#
#     return vel, I, mu, mu_err, I_m, f



fig, ax = pl.subplot_mosaic(
    [["a", "a"],
     ["a", "a"],
     ["b", "c"]],
    layout="constrained", figsize=(7. / 2 * 1.5, 3.8 * 1.5))

meas_list = [
    {
        "path": "Exp_Shapiro_steps_f_90_Hz_I_m_0.0.csv",
        "mod_amp": 0.,
        "color": "#2e2e41",
    },
    {
        "path": "Exp_Shapiro_steps_f_90_Hz_I_m_1.0.csv",
        "mod_amp": 0.7,
        "color": "navy",
    },
    {
        "path": "Exp_Shapiro_steps_f_90_Hz_I_m_1.6.csv",
        "mod_amp": 1.2,
        "color": "#fb6d10"
    },


]

n = len(meas_list)
colors = [RPTU_COLORS[c] for c in ["nacht", "schiefer", "mango", ]]


##### A ######
for i, meas in enumerate(meas_list):

    I_c = current(1, v_c=1)  # get the critical current from the func keyword parameter
    v_max = 0.7

    vel, I, mu, mu_err, I_m, f = get_meas_data(meas["path"], path_data_repo = path_data_repo)
    style = get_style(colors[i])

    if I_m == 0:
        ax["a"].errorbar(I[vel <= v_max] / I_c, mu[vel <= v_max], mu_err[vel <= v_max], **style,
                         label=fr'$I_\mathrm{{m}} = {0}$')
    else:
        ax["a"].errorbar(I[vel <= v_max] / I_c, mu[vel <= v_max], mu_err[vel <= v_max], **style,
                         label=fr'$I_\mathrm{{m}} = {I_m / I_c:.1f}$ $I_\mathrm{{c}}$')



ax["a"].axhline(y=90, color='dimgrey', ls='--')
ax["a"].axhline(y=180, color='dimgrey', ls='--')

ax["a"].annotate(text='$1 \cdot f_\mathrm{m}$', xy=(0, 90), verticalalignment='center', fontsize=12,
                 bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'dimgrey', 'ls': '-', 'lw': 1.5}, zorder=5)

ax["a"].annotate(text='$2 \cdot f_\mathrm{m}$', xy=(0, 2 * 90), verticalalignment='center', fontsize=12,
                 bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'dimgrey', 'ls': '-', 'lw': 1.5}, zorder=5)

#  label = r'$n \cdot f_\mathrm{m}$'


ax["a"].set_ylim(top=264)
ax["a"].get_yaxis().set_ticks([0, 100, 200])
ax["a"].get_yaxis().set_ticklabels([0, 100, 200], fontsize=12)

ax["a"].get_xaxis().set_ticks([0, 0.5, 1, 1.5])
ax["a"].get_xaxis().set_ticklabels([0, 0.5, 1, 1.5], fontsize=12)



##### B ############


data = {
    "path": "Exp_Shapiro_steps_f_90_Hz_I_m_1.0.csv",
    "mod_amp": 0.7,
}

vel, I, mu, mu_err, I_m, f = get_meas_data(data["path"], path_data_repo = path_data_repo)

ax["b"].errorbar(I[vel <= v_max] / I_c, mu[vel <= v_max], mu_err[vel <= v_max], **get_style(RPTU_COLORS["schiefer"]),
                 label=fr'$I_\mathrm{{m}} = {I_m / I_c:.1f}$ $I_\mathrm{{c}}$')

ax["b"].axhline(y=90, color='dimgrey', ls='--')
ax["b"].axhline(y=180, color='dimgrey', ls='--')

ax["b"].annotate(text='$1 \cdot f_\mathrm{m}$', xy=(0, 90), verticalalignment='center', fontsize=8,
                 bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'dimgrey', 'ls': '-', 'lw': 1.5}, zorder=5)



# plot b experiment theory comparison
tab_b = pd.read_csv(join(path_data_repo, "Sim_Shapiro_steps_f_90_Hz_I_m_1.0.csv"))
vel = tab_b["I/I_c"].to_numpy()
chem_pot_data = tab_b["chem_pot"].to_numpy()
I_m = tab_b["I_m"].to_numpy()[0]
ax['b'].errorbar(vel, chem_pot_data, **get_style(RPTU_COLORS["himbeere"], ls='-', marker='.'),
                 label=fr"Sim. {I_m:.1f} $I_c$")



##### C ######
# Simulation data

sim_list_names = ["Sim_Shapiro_steps_f_90_Hz_I_m_0.0.csv",
                  "Sim_Shapiro_steps_f_90_Hz_I_m_0.5.csv",
                  "Sim_Shapiro_steps_f_90_Hz_I_m_0.7.csv",
                  "Sim_Shapiro_steps_f_90_Hz_I_m_1.0.csv",
                  "Sim_Shapiro_steps_f_90_Hz_I_m_1.4.csv",]

sim_list_mod_amps = np.asarray([0, 0.4, 0.5, 0.75, 1]) * 1e-3  # modulation amplitudes in µm

n = len(sim_list_names)
colors = rptu_centered(np.linspace(0.2, 0.8, n))



# plot c theory
for i, name in enumerate(sim_list_names):

    table = pd.read_csv(join(path_data_repo, name))
    vel = table["I/I_c"].to_numpy()
    chem_pot_data = table["chem_pot"].to_numpy()
    I_m = table["I_m"].to_numpy()[0]

    ax['c'].errorbar(vel, chem_pot_data, **get_style(colors[i], ls='-', marker='.'),
                     label=fr"Sim. {I_m:.1f} $I_c$")

ax["c"].axhline(y=90, color='dimgrey', ls='--')
ax["c"].axhline(y=2 * 90, color='dimgrey', ls='--')
ax["c"].axhline(y=3 * 90, color='dimgrey', ls='--')




ax["a"].set_xlim(right=1.7)
ax["a"].set_xlabel(r"Current $I \, / \, I_\mathrm{c}$")
ax["a"].set_ylabel(r" $\Delta \mu \,  / \, h$ [Hz] ")
ax["a"].legend(loc='upper left')

ax["b"].set_xlim(right=1.65)
ax["b"].legend(loc='upper left', fontsize=8)
ax["b"].set_xlabel(r" $I  \,  / \, I_\mathrm{c}$")
ax["b"].set_ylabel("$\Delta \mu \,  / \, h$ [Hz]")

ax["c"].set_xlim(right=1.85)
ax["c"].legend(loc='upper left', fontsize=6)
ax["c"].set_xlabel(r" $I  \,  / \, I_\mathrm{c}$")
ax["c"].set_ylabel(r"$\Delta \mu \,  / \, h$ [Hz]")



fig.tight_layout()
pl.show()
pl.close()
