"""
Helper functions for evaluations
"""
import os.path, glob
from os.path import join
import json
import pandas as pd
import numpy as np
import matplotlib.pyplot as pl

# Ref chemical potential
chem_pot_ref = 3354

""" Rescaling factor for the velocity from Hz/ms to mm/s
"""
scaling = 3.4 * 1e-6 # in µm / MHz in atom plane

def mu_N(N):
    """
    func to calc the chemical potential for the sample geometry depending on atom number
    trap parameter w = 2*pi*[1.6,252,250] Hz
    :param N: atom number
    :return: µ(N) in Hz
    """
    return 4.619 * N ** (1 / 2)



def mu(z, N):
    """
    func to calculate the chemical potential difference from imbalance,
    depending on atom number
    :param z: imbalance
    :param N: atom number
    :return: delta µ in Hz
    """
    return mu_N(N * (1 + z)) - mu_N(N * (1 - z))


def current(v,v_c = 0.42, I_c=192e3):
    """
    Function to resale the velocity to the current.
    NOTE: The I_c here is measured for bar_pow = 0.45µ; each barrier height has its own I_c
    The same holds for v_c
    :param v: velocity in mm/s
    :param v_c: critical velocity in mm/s
    :param I_c: critical current in 1/s
    :return: current in 1/s
    """
    I = v * I_c / v_c
    return I

def get_mod_amp(mod_amp, f):
    """
    function to get the modulation amplitude as current from a modulation amplitude in µm
    :param mod_amp: modulation amplitude in µm
    :param f: modulation frequency
    :return: mod amp in current
    """
    mod_velocity = 2 * np.pi * f * mod_amp * 1e-3
    # mod amp in µm need to be multiplied with 1e-3 to get mod_amp in mm, mod_velocity in [mm/s]
    # since the critical velocity is in mm/s

    # use the current function to convert from velocity in current
    I_m = current(mod_velocity)

    return I_m



def get_meas_data(meas, path_data_repo = None,mod_amp_return = False, meas_dict = None):
    """
    function to read the csv table of the measurement Shapiro data tables
    :returns
    vel: velocity n mm/s,
    I: current in 1/s,
    mu, Chemical potential difference in Hz
    mu_err, Error of the Chemical potential difference
    I_m Modulation current in I_m / I_c,
    f modulation frequency in Hz
    """
    #print(meas)
    path_exp = os.path.join(path_data_repo, meas)

    # Read Experiment Data
    data = pd.read_csv(path_exp)
    # Calculate chemical potential difference
    data["chem_pot"] = mu_N(data["N"] * (1 + data['delta_z'])) - mu_N(
        data["N"] * (1 - data['delta_z']))
    # get frequency
    f = data["f"].to_numpy()[0]
    # get modulation amplitude
    I_m = data["I_m"].to_numpy()[0] # I_m / I_c
    # get critical current from func keyword
    I_c = current(1, v_c=1)

    if mod_amp_return:
        try:
            mod_amp = data["modulation_amplitude"].to_numpy()[0] *1e-6 *  3.4 # returns modulation amplitude in µm
        except KeyError:
            print("Key Error: modulation_amplitude not in table.keys() \n Use fall back dict value")
            try:
                mod_amp = meas_dict['mod_amp']
            except KeyError:
                print("No fallback value found. Proceed without value. ")
                mod_amp = None

    data = data[["delta_z", "vel", "chem_pot"]]
    # calculate the average of the data for the unique velocities
    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)

    if mod_amp_return:
        return vel, I, mu, mu_err, I_m, f, mod_amp, z, z_err
    else:
        return vel, I, mu, mu_err, I_m, f



### sigmoid functions for fitting the Shapiro steps

def sigmoid_1(x, L, k1, k2, k3, x1, x2, x3, off):
    f = L / (1 + np.exp(-k1 * (x - x1))) + off
    return f


def sigmoid_2(x, L, k1, k2, k3, x1, x2, x3, off):
    f = L / (1 + np.exp(-k1 * (x - x1))) + L / (
            1 + np.exp(-k2 * (x - (x2 + x1)))) + off
    return f


def sigmoid_3(x, L, k1, k2, k3, x1, x2, x3, off):
    f = L / (1 + np.exp(-k1 * (x - x1))) + L / (1 + np.exp(-k2 * (x - (x2 + x1)))) + L / (
            1 + np.exp(-k3 * (x - (x3 + x2 + x1)))) + off
    return f




def load_image_dict(dir_path, seqTimes=None):
    """loads an image from a list of timestemps, {key seqtimestemp : value path} """
    all_h5 = glob.glob(join(dir_path, '*.h5'))
    image_dict = {}
    for val in all_h5:
        seq = val.split('/')[-1].split('_')[-1].split('.')[0]
        if seqTimes is not None:
            if np.isin(seq, seqTimes):
                image_dict[int(seq)] = val
            else:
                pass
        else:
            image_dict[int(seq)] = val
    return image_dict


def conductance_fit(x, v_c, G, off):
    # off = 0
    result = off * np.ones_like(x)
    result[x >= v_c] += np.sqrt(x[x >= v_c] ** 2 - v_c ** 2) / G
    return result

def resistance_fit(x, R, off, v_c):
    # off = 0
    result = off * np.ones_like(x)
    result[x >= v_c] += np.sqrt(x[x >= v_c] ** 2 - v_c ** 2) * R
    return result