"""
File to evaluate the Shapiro curves we have measured to get the frequency dependence of the step height shown in Fig 2 A.
"""


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

import pandas as pd
from scipy.optimize import curve_fit
import numpy as np
import matplotlib.pyplot as pl

from helper import *
from style import *


meas_list = [

    {
        "path": "Exp_Shapiro_steps_f_121_Hz_I_m_1.5.csv",
        "mod_amp": 0.85,
        "mod_freq": 120,
        "p0": [65, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "forestgreen",
        "steps": 3,
        'v_max': 0.8
        },
    {
        "path":"Exp_Shapiro_steps_f_121_Hz_I_m_0.5.csv" ,
        "mod_amp": 0.3,
        "mod_freq": 120,
        "p0": [65, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "forestgreen",
        "steps": 3,
        'v_max': 1
    },
    {
        "path": "Exp_Shapiro_steps_f_121_Hz_I_m_0.9.csv",
        "mod_amp": 0.5,
        "mod_freq": 120,
        "p0": [65, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "forestgreen",
        "steps": 3,
        'v_max': 1
    },

    {
        "path":"Exp_Shapiro_steps_f_61_Hz_I_m_0.9.csv",
            "mod_amp": 1.,
        "mod_freq": 60,
        "color": "navy",
        "p0": [60, 16, 25, 50, 0.2, 0.2, 0.55, -5],
        "steps": 3,
        'v_max': 0.6
        },

    {
        "path":"Exp_Shapiro_steps_f_61_Hz_I_m_1.4.csv",
        "mod_amp": 1.5,
        "mod_freq": 60,
        "color": "navy",
        "p0": [60, 50, 50, 50, 0.1, 0.2, 0.55, -5],
        "steps": 3,
        'v_max': 0.8
        },

    {
        "path":"Exp_Shapiro_steps_f_61_Hz_I_m_1.5.csv",
        "mod_amp": 1.7,
        "mod_freq": 60,
        "color": "navy",
        "p0": [60, 50, 50, 50, 0.1, 0.2, 0.55, -5],
        "steps": 3,
        'v_max': 0.8
        },



    {
        "path":"Exp_Shapiro_steps_f_91_Hz_I_m_1.4.csv",
        "mod_amp": 1,
        "mod_freq": 90,
        "p0": [90, 50, 100, 100, 0.05, 0.25, 0.55, 5],
        "color": "darkgreen",
        "steps": 3,
        'v_max': 0.8
    },
    {
        "path":"Exp_Shapiro_steps_f_91_Hz_I_m_0.7.csv",
        "mod_amp": 0.5,
        "mod_freq": 90,
        "p0": [80, 40, 45, 40, 0.3, 0.2, 0.2, 5],
        "color": "darkgreen",
        "steps": 3,
        'v_max': 0.8
    },
    {
        "path":"Exp_Shapiro_steps_f_91_Hz_I_m_1.0.csv",
        "mod_amp": 0.7,
        "mod_freq": 90,
        "p0": [65, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color":"darkgreen",
        "steps":2,
        'v_max': 0.8
      },
    {
            "path":"Exp_Shapiro_steps_f_91_Hz_I_m_0.8.csv",
            "mod_amp": 0.6,
            "mod_freq": 90,
            "p0": [90, 50, 100, 100, 0.05, 0.25, 0.55, 5],
            "color": "darkgreen",
            "steps": 3,
            'v_max': 0.8
        },
    {
        "path":"Exp_Shapiro_steps_f_136_Hz_I_m_1.0.csv",
        "mod_amp": 0.5,
        "mod_freq": 135,
        "p0": [135, 40, 45, 40, 0.3, 0.2, 0.7, 5],
        "color": "orange",
        "steps": 3,
        'v_max': 1
    },
    #
    {
        "path":"Exp_Shapiro_steps_f_136_Hz_I_m_1.7.csv",
        "mod_amp": 0.85,
        "mod_freq": 135,
        "p0": [140, 40, 45, 40, 0.1, 0.3, 0.3, 5],
        "color": "orange",
        "steps": 3,
        'v_max': 1.
    },
    {
        "path":"Exp_Shapiro_steps_f_76_Hz_I_m_0.4.csv",
        "mod_amp": 0.35,
        "mod_freq": 75,
        "p0": [100, 100, 100, 100, 0.2, 0.1, 0.3, 5],
        "color": "purple",
        "steps": 3,
        'v_max': 0.7
            },
    {
        "path":"Exp_Shapiro_steps_f_76_Hz_I_m_0.9.csv",
        "mod_amp": 0.8,
        "mod_freq": 75,
        "p0": [75, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "purple",
        "steps": 3,
        'v_max': 0.8
    },
    {
        "path":"Exp_Shapiro_steps_f_76_Hz_I_m_1.1.csv",
        "mod_amp": 1,
        "mod_freq": 75,
        "p0": [75, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "purple",
        "steps": 3,
        'v_max': 0.8
            },


    {
        "path":"Exp_Shapiro_steps_f_106_Hz_I_m_1.1.csv",
        "mod_amp": 0.7,
        "mod_freq": 105,
        "p0": [105, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "#e37239",
        "steps": 3,
        'v_max': 0.8
    },
    {
        "path":"Exp_Shapiro_steps_f_106_Hz_I_m_1.6.csv",
        "mod_amp": 1,
        "mod_freq": 105,
        "p0": [75, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "#e37239",
        "steps": 3,
        'v_max': 0.8
    },



    {
        "path": "Exp_Shapiro_steps_f_106_Hz_I_m_0.8.csv",
        "mod_amp": 0.5,
        "mod_freq": 105,
        "p0": [75, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "#e37239",
        "steps": 3,
        'v_max': 0.8
    },

    {
        "path":"Exp_Shapiro_steps_f_152_Hz_I_m_1.1.csv",
        "mod_amp": 0.5,
        "mod_freq": 150,
        "p0": [150, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "royalblue",
        "steps": 3,
        'v_max': 1
    },

    {
        "path": "Exp_Shapiro_steps_f_152_Hz_I_m_1.6.csv",
        "mod_amp": 0.7,
        "mod_freq": 150,
        "p0": [45, 50, 100, 100, 0.05, 0.25, 0.55, 5],
        "color": "royalblue",
        "steps": 2,
        'v_max': 1
    },

    {
        "path": "Exp_Shapiro_steps_f_152_Hz_I_m_1.6.csv",
        "mod_amp": 0.7,
        "mod_freq": 150,
        "p0": [150, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "royalblue",
        "steps": 3,
        'v_max': 1
    },

    {
        "path":"Exp_Shapiro_steps_f_182_Hz_I_m_1.4.csv",
        "mod_amp": 0.5,
        "mod_freq": 180,
        "p0": [180, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "steelblue",
        "steps": 2,
        'v_max': 1
    },

    {
        "path":"Exp_Shapiro_steps_f_182_Hz_I_m_1.0.csv",
        "mod_amp": 0.35,
        "mod_freq": 180,
        "p0": [200, 100, 100, 100, 0.4, 0.4, 0.55, 5],
        "color": "steelblue",
        "steps": 2,
        'v_max': 1
    },


    {
        "path":"Exp_Shapiro_steps_f_91_Hz_I_m_1.2.csv",
        "mod_amp": 0.9,
        "mod_freq": 90,
        "p0": [100, 100, 100, 100, 0.2, 0.3, 0.25, 5],
        "color": "steelblue",
        "steps": 3,
        'v_max': 0.8
    },

    {
        "path":"Exp_Shapiro_steps_f_167_Hz_I_m_1.7.csv",
        "mod_amp": 0.7,
        "mod_freq": 165,
        "p0": [45, 50, 100, 100, 0.05, 0.25, 0.55, 5],
        "color": "darkorange",
        "steps": 2,
        'v_max': 1
    },

    {
        "path":"Exp_Shapiro_steps_f_167_Hz_I_m_1.2.csv",
        "mod_amp": 0.5,
        "mod_freq": 165,
        "p0": [165, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "darkorange",
        "steps": 3,
        'v_max': 1
    },

    {
        "path":"Exp_Shapiro_steps_f_167_Hz_I_m_0.7.csv",
        "mod_amp": 0.3,
        "mod_freq": 165,
        "p0": [165, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "darkorange",
        "steps": 2,
        'v_max': 1
    },

    {
        "path":"Exp_Shapiro_steps_f_91_Hz_I_m_1.0.csv",
        "mod_amp": 0.7,
        "mod_freq": 90,
        "p0": [60, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "darkorange",
        "steps": 3,
        'v_max': 1
    },


    {
        "path":"Exp_Shapiro_steps_f_182_Hz_I_m_1.9.csv",
        "mod_amp": 0.7,
        "mod_freq": 180,
        "p0": [180, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "firebrick",
        "steps": 2,
        'v_max': 1,
        'v_plot': 0.6,
        'plot': False
    },
    {
        "path":"Exp_Shapiro_steps_f_136_Hz_I_m_1.4.csv",
        "mod_amp": 0.7,
        "mod_freq": 135,
        "p0": [135, 100, 100, 100, 0.108, 0.25, 0.55, 5],
        "color": "darkgreen",
        "steps": 3,
        'v_max': 1,
        'v_plot': 0.4,
        'plot': True
    },
    {
        "path":"Exp_Shapiro_steps_f_91_Hz_I_m_1.6.csv",
        "mod_amp": 1.2,
        "mod_freq": 90,
        "p0":[90, 50, 100, 100, 0.05, 0.25, 0.55, 5],
        "color": "orange",
        "steps": 3,
        'v_max': 0.65,
        'v_plot': 0.35,
        'plot': True
    },

    {
        "path":"Exp_Shapiro_steps_f_61_Hz_I_m_1.1.csv",
        "mod_amp": 1.2,
        "mod_freq": 60,
        "color": "navy",
        "p0": [60,50,50,50,0.15,0.2,0.2,-10],
        "steps": 3,
        'v_max': 0.65,
        'v_plot': 0.32,
        'plot': True
        },

]


path_data_repo = "../Fig_2/shapiro_curves"

fig, ax = pl.subplots(nrows=1, ncols=1,figsize=(3,4))
axs = ax.flatten()

result = []

n = 4
colors = cmaps.torch(np.linspace(0.2,0.7,n))
colors = colors[::-1]
i=0

for meas in meas_list:
    # get data
    vel, I, mu, mu_err, I_m, f = get_meas_data(meas["path"],path_data_repo = path_data_repo)

    v_max = meas['v_max']

    I = current(vel, v_c=0.42)

    # choose function depending on the visible steps
    try:
        if meas["steps"] == 2:
            sigmoid = sigmoid_2
        elif meas["steps"] ==1:
            sigmoid = sigmoid_1
        else:
            sigmoid=sigmoid_3

        popt, pcov = curve_fit(sigmoid,vel[vel <= v_max],mu[vel <= v_max],p0=meas["p0"],bounds = (-10,np.inf))

        v = np.arange(0,vel[-1],0.001)
        p_err = np.sqrt(np.diag(pcov))

        # def sigmoid(x, L, k1, k2, k3 , x1, x2, x3):
        # get step height from fit
        L = popt[0]
        L_err = p_err[0]

        # put it in list
        result.append([f, meas["mod_amp"], L, L_err])

        # uncomment to plot single curves
        # fig = pl.figure(figsize=(9, 6))
        # pl.errorbar(vel, mu, mu_err, **get_style(meas['color']),
        #             label=f'$A_m = {meas["mod_amp"]:.1f}\mathrm{{µm}}$, f= {meas["mod_freq"]} Hz')
        # pl.plot(v, sigmoid(v, *popt), color=meas["color"])
        # pl.legend()
        # pl.savefig(f'{meas["path"]}.png')
        #pl.show()
    except RuntimeError:
        print('Fit Failed!')

# save result in data frame
result = np.asarray(result)
df  = pd.DataFrame(result, columns=['mod_freq','mod_amp', 'L', 'L_err'])


def lin(x,m):
    return m*x



df = df[["mod_freq", "L" ]]
grouped = df.groupby(by="mod_freq").agg(['mean', 'std'])
mod_freq = grouped.index.to_numpy()

popt , pcov = curve_fit(lin, df["mod_freq"], df["L"])

# plot µ=f slope
x = np.arange(0,200,0.01)
axs[0].plot(x,lin(x,1), color = '#e88607', label= fr'$\mu = f_\mathrm{{m}}$')

data_dict = {'mod_freq':mod_freq,
             'step_height':grouped["L", "mean"],
             'step_height_error':grouped["L", "std"]}

# save result data frame
df = pd.DataFrame(data_dict)
#df.to_csv(os.path.join(GROUP_STORAGE_PATH,"Projects/Shapiro/Shapiro_steps/frequency_dependence/step_height_data.csv"),index = False)

axs[0].errorbar(mod_freq,grouped["L", "mean"],grouped["L", "std"], **get_style(color='#12223b'), markersize = 7,markeredgewidth=1.8)
axs[0].set_ylabel(r"Chemical Potential $\Delta \mu$ [Hz]")
axs[0].set_xlabel(r"Modulation Frequency $f_\mathrm{m}$ [Hz]")
axs[0].legend(loc = 'upper left')
pl.tight_layout()

pl.show()