import os.path, glob
from os.path import join
import json

import pandas as pd
import numpy as np
import matplotlib.pyplot as pl
import h5py
from style import *
from helper import *

path_data_repo = "../Fig_3"

# data dict DC measurements Experiments
meas_list_dc = [
    {
        "path": "Exp_DC_density_time_evolution_I_0.48.h5",
        "mod_amp": 0.0,
        "mod_freq": 90,
        "velocity": 0.2,
        'label': r'$\mathrm{I}$'
    },

    {
        "path": "Exp_DC_density_time_evolution_I_0.95.h5",
        "mod_amp": 0.0,
        "mod_freq": 90,
        "velocity": 0.4,
        'label': r'$\mathrm{II}$'
    },

    {
        "path": "Exp_DC_density_time_evolution_I_1.43.h5",
        "mod_amp": 0.0,
        "mod_freq": 90,
        "velocity": 0.6,
        'label': r'$\mathrm{III}$'
    },
]

# data dict ac_measuremnts Experiment
meas_list_ac = [

    {
        "path": "Exp_AC_density_time_evolution_I_0.24.h5",
        "mod_amp": 0.6,
        "mod_freq": 90,
        "velocity": 0.1,
        'label': r'$\mathrm{I}$'
        },
    {
        "path": "Exp_AC_density_time_evolution_I_0.88.h5",
        "mod_amp": 0.6,
        "mod_freq": 90,
        "velocity": 0.37,
        'label': r'$\mathrm{II}$'
    } ,
    {
        "path": "Exp_AC_density_time_evolution_I_1.31.h5",
        "mod_amp": 0.6,
        "mod_freq": 90,
        "velocity": 0.55,
        'label': r'$\mathrm{III}$'
    },
]

colormap = cmaps.vik


######## DC measurements ##########


set_plot_style()
fig, ax = pl.subplots(nrows=1, ncols=int(len(meas_list_dc)), sharey= True, figsize = (5 * 1.61,3),dpi = 300)
axs = ax.flatten()

# maximal density
v_max = 1000
for i, meas in enumerate(meas_list_dc):
    path = os.path.join(path_data_repo, meas["path"])
    # open files
    h5_file = h5py.File(path, 'r')
    image = h5_file['image']

    # get meta_data
    y_max = h5_file.attrs['t_max']
    x_max = h5_file.attrs['conversion_to_um'] * image.shape[1]
    pixel_size  = h5_file.attrs['pixel_size']
    magnification = h5_file.attrs['magnification']

    # transform the atomic density back to atom space
    image = np.asarray(image) / pixel_size * magnification  # [N/µm] ; Pixel size = 6.45 µm/Pixel

    # plot image
    im = axs[i].imshow(image, cmap=colormap , aspect='auto', origin='lower', extent=(0, x_max, 0, y_max), vmin=-v_max, vmax=v_max)

    ticks = np.linspace(0,60,4,dtype=int)
    axs[i].set_xlabel('x [µm]', fontsize = 12)
    axs[i].text(x = 38, y = 36, s=meas['label'],ha = 'center', verticalalignment='center', fontsize=10,
                     bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)
    axs[i].set_xlim(0,65)


# comment in to get a colorbar
#cbar = pl.colorbar(im, ax = axs[i-1])
#ax[i-1].cax.colorbar(im)
#cbar.ax.set_ylabel(r'$\Delta n$ [1/µm]', fontsize = 12)
#cbar.ax.set_yticks([100,0,-100],direction = "out")
axs[0].set_ylabel("Time [ms]",fontsize = 12)

fig.tight_layout()
#pl.savefig('EXP_DC_Density_modulation_overview.png')




######## AC measurements ######


#set_plot_style()

fig, ax = pl.subplots(nrows=1, ncols=int(len(meas_list_ac)), sharey= True, figsize = (5 * 1.61,3),dpi = 300)
axs = ax.flatten()

for i, meas in enumerate(meas_list_ac):
    path = os.path.join(path_data_repo, meas["path"])
    # open files
    h5_file = h5py.File(path, 'r')
    image = h5_file['image']

    # get meta_data
    y_max = h5_file.attrs['t_max']
    x_max = h5_file.attrs['conversion_to_um'] * image.shape[1]
    pixel_size  = h5_file.attrs['pixel_size']
    magnification = h5_file.attrs['magnification']

    # transform the atomic density back to atom space
    image = np.asarray(image) / pixel_size * magnification  # [N/µm] ; Pixel size = 6.45 µm/Pixel

    # plot image
    im = axs[i].imshow(image, cmap=colormap, aspect='auto', origin='lower', extent=(0, x_max, 0, y_max), vmin=-v_max, vmax=v_max,)

    ticks = np.linspace(0,60,4,dtype=int)
    axs[i].set_xlabel('x [µm]',fontsize = 12)
    axs[i].text(x = 38, y = 36, s=meas['label'],ha = 'center', verticalalignment='center', fontsize=10,
                     bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)
    axs[i].set_xlim(0, 65)

# comment in to get colorbar
#cbar = pl.colorbar(im, ax = axs[i])
#ax[i-1].cax.colorbar(im)
#cbar.ax.set_ylabel(r'$\Delta n$ [1/µm]',fontsize =12 )
#cbar.ax.set_yticks([100,0,-100],direction = "out")
axs[0].set_ylabel("Time [ms]",fontsize = 12)

fig.tight_layout()
#pl.savefig('EXP_AC_Density_modulation_overview.png')


########################################################################################################
########################################################################################################

# Simulation data


data_dict_Sim = [{"name":"Sim_AC_density_time_evolution_I_0.50.h5",
              'label': r'$\mathrm{I}$',
              'I':0.5},#I_dc/I_c

                {"name":"Sim_AC_density_time_evolution_I_0.93.h5",
                 'label': r'$\mathrm{II}$',
                 'I':0.93},

                {"name":"Sim_AC_density_time_evolution_I_1.50.h5",
                 'label': r'$\mathrm{III}$',
                 'I':1.5}
             ]


set_plot_style()
fig, ax = pl.subplots(nrows=1, ncols=int(len(data_dict_Sim)), sharey= True, figsize = (5 * 1.61,3),dpi = 300)
axs = ax.flatten()

v_max = 1000 # max density

for i, data in enumerate(data_dict_Sim):
    # open image
    path = os.path.join(path_data_repo, data["name"])
    # open files
    h5_file = h5py.File(path, 'r')
    image = h5_file['image']

    # get meta_data
    x_step_size = h5_file.attrs['x_step_size']
    t_step_size = h5_file.attrs['t_step_size']
    x_grid = h5_file.attrs['x_grid']
    t_grid = h5_file.attrs['t_grid']
    I = h5_file.attrs['I']

    #
    image = np.asarray(image) / x_step_size # to get n [1 / µm]

    x_max = x_grid * x_step_size
    t_max = t_grid * t_step_size

    x_min = 0
    x_max = np.shape(image)[1] * x_step_size

    extent = (0, x_max, 0, t_max)
    # plot image
    im = axs[i].imshow(image, cmap=cmcrameri.cm.vik , aspect='auto', origin='lower', extent=extent, vmin=-v_max, vmax=v_max)

    ticks = np.linspace(0,60,4,dtype=int)
    axs[i].get_xaxis().set_ticklabels(ticks)
    axs[i].get_xaxis().set_ticks(ticks)
    axs[i].set_xlabel('x [µm]', fontsize = 12)

    axs[i].text(x = 30, y = 36, s=data['label'],ha = 'center', verticalalignment='center', fontsize=10,
                     bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)

    axs[i].set_xlim(0,65)

ax[0].set_ylabel("Time [ms]", fontsize = 12)

# un comment if you need a colorbar
# cbar = pl.colorbar(im, ax = axs[i])
# cbar.ax.set_ylabel(r'$\Delta n$ [1/µm]',fontsize =12 )

fig.tight_layout()
#pl.savefig('SIM_Density_modulation_overview.png')

##################################################################################
##################################################################################

#  plot shapiro curves

# Experiment

meas_list = [
    {
        "path": "Exp_AC_Shapiro_steps_I_m_0.82.csv",
        "mod_amp": 0.6,
        "mod_freq": 90,
        "p0": [90, 50, 100, 100, 0.05, 0.25, 0.55, 5],
        "steps": 3,
        'v_max': 0.8
    },

    {
        "path": "Exp_DC_Shapiro_steps_I_m_0.csv",
        "mod_amp": 0.,
        "mod_freq": 0,
        'v_max': 0.8
    },

]


#####  DC ##########

set_plot_style()

pl.figure(figsize=(3 * 1.2, 3))
ax = pl.gca()

# get data
vel, I, mu, mu_err, I_m, f = get_meas_data(meas_list[1]['path'],path_data_repo=path_data_repo)

I_c = current(1, v_c=1)

# plot data
ax.errorbar(I[vel <= 0.8] / I_c, mu[vel <= 0.8], mu_err[vel <= 0.8], **get_style(RPTU_COLORS['nacht']),label=fr'$I_\mathrm{{m}} = {0}$')

# mark the current values where the density time evolution was measured, see meas_list_dc
velocity_values = np.array([0.2, 0.4, 0.6])
I_values = current(velocity_values)

ax.vlines(I_values[0] / I_c, 280, 0, color='orange', ls='--')
ax.vlines(I_values[1] / I_c, 280, 10, color='orange', ls='--')
ax.vlines(I_values[2] / I_c, 280, 175, color='orange', ls='--')

ax.errorbar(I_values[0] / I_c, 0, color='orange', ls='none', marker='o')
ax.errorbar(I_values[1] / I_c, 10, color='orange', ls='none', marker='o')
ax.errorbar(I_values[2] / I_c, 175, color='orange', ls='none', marker='o')

ax.annotate(text=r'$\mathrm{I}$', xy=(I_values[0] / I_c, 270), ha='center', verticalalignment='center', fontsize=10,
            bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)

ax.annotate(text=r'$\mathrm{II}$', xy=(I_values[1] / I_c, 270), ha='center', verticalalignment='center', fontsize=10,
            bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)

ax.annotate(text=r'$\mathrm{III}$', xy=(I_values[2] / I_c, 270), ha='center', verticalalignment='center', fontsize=10,
            bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)

ax.set_xlabel(r"Current $I/I_c$", fontsize=12)
ax.set_ylabel("$\Delta \mu \, / \, h$ [Hz]", fontsize=12)

pl.tight_layout()
#pl.savefig('Shapiro_curve_dc.png')


########## AC ############

pl.figure(figsize=(3 * 1.2, 3))
ax = pl.gca()

# get data
vel, I, mu, mu_err, I_m, f = get_meas_data(meas_list[0]['path'],path_data_repo=path_data_repo)

# plot data
ax.errorbar(I[vel <= 0.8] / I_c, mu[vel <= 0.8], mu_err[vel <= 0.8], **get_style(RPTU_COLORS['schiefer']),
            label=fr'$I_\mathrm{{m}} = {I_m / I_c:.2f}$ $I/I_c$')

# mark the current values where the density time evolution was measured, see meas_list_ac
velocity_values = np.array([0.1, 0.37, 0.55])
I_values = current(velocity_values)

ax.vlines(I_values[0] / I_c, -60, 0, color='orange', ls='--')
ax.vlines(I_values[1] / I_c, -30, 95, color='orange', ls='--')
ax.vlines(I_values[2] / I_c, -30, 190, color='orange', ls='--')

ax.errorbar(I_values[0] / I_c, 0, color='orange', ls='none', marker='o')
ax.errorbar(I_values[1] / I_c, 95, color='orange', ls='none', marker='o')
ax.errorbar(I_values[2] / I_c, 190, color='orange', ls='none', marker='o')

ax.annotate(text=r'$\mathrm{I}$', xy=(I_values[0] / I_c, -40), ha='center', verticalalignment='center', fontsize=10,
            bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)

ax.annotate(text=r'$\mathrm{II}$', xy=(I_values[1] / I_c, -40), ha='center', verticalalignment='center', fontsize=10,
            bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)

ax.annotate(text=r'$\mathrm{III}$', xy=(I_values[2] / I_c, -40), ha='center', verticalalignment='center', fontsize=10,
            bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)


ax.set_xlabel(r"Current $I/I_c$", fontsize=12)
ax.set_ylabel("$\Delta \mu \, / \, h$ [Hz]", fontsize=12)

pl.tight_layout()
#pl.savefig('Shapiro_curve_ac.png')


######  Simulation #############


pl.figure(figsize=(3 * 1.2, 3))
ax = pl.gca()

################ AC case

table = pd.read_csv(join(path_data_repo,"Sim_AC_Shapiro_steps_f_90_Hz_I_m_1.9.csv"))

I = table["I/I_c"].to_numpy()  # mm/s
mu = table["chem_pot"].to_numpy()

I_m = table["I_m"].to_numpy()[0]


color = RPTU_COLORS['himbeere']

ax.errorbar(I, mu, **get_style(color, errorbar=False, ls=' ', marker='v'),
            label=fr'Sim. $I_\mathrm{{m}} = {I_m}$ $I/I_c$')

#  mark the current values where the density time evolution was measured, see data_dict_Sim
velocity_values = np.array([0.52, 0.95, 1.3])
I_values = current(velocity_values)

ax.vlines(velocity_values[0], -60, 5, color='orange', ls='--')
ax.vlines(velocity_values[1], -30, 95, color='orange', ls='--')
ax.vlines(velocity_values[2], -30, 195, color='orange', ls='--')

ax.errorbar(velocity_values[0], 5, color='orange', ls='none', marker='o')
ax.errorbar(velocity_values[1], 95, color='orange', ls='none', marker='o')
ax.errorbar(velocity_values[2], 195, color='orange', ls='none', marker='o')

ax.annotate(text=r'$\mathrm{I}$', xy=(velocity_values[0], -40), ha='center', verticalalignment='center', fontsize=10,
            bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)

ax.annotate(text=r'$\mathrm{II}$', xy=(velocity_values[1], -40), ha='center', verticalalignment='center', fontsize=10,
            bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)

ax.annotate(text=r'$\mathrm{III}$', xy=(velocity_values[2], -40), ha='center', verticalalignment='center', fontsize=10,
            bbox={'boxstyle': 'round', 'fc': "white", 'ec': 'orange', 'ls': '-', 'lw': 1.5}, zorder=5)

ax.set_xlabel(r"Current $I/I_c$", fontsize=12)
ax.set_ylabel("$\Delta \mu \, / \, h$ [Hz]", fontsize=12)

pl.tight_layout()
#pl.savefig('Shapiro_curve_Sim.png')

pl.show()