import numpy as np
import matplotlib.pyplot as pl
import pandas as pd
import os, glob
from os.path import join

from scipy.optimize import curve_fit

from scripts.helper import *
from scripts.style import *

path_data_repo = "../data_supplementary"



def lin(x,x0,m,off):
    return m * (x - x0) + off


set_plot_style(PYPLOT_STYLE_WHITE)
fig = pl.figure(figsize=(3.5 * 1.61, 3.5))


print(os.path.join(path_data_repo,'barrier_height_calibration.csv'))
data = pd.read_csv(os.path.join(path_data_repo,'barrier_height_calibration.csv'))


bar_pow = data['bar_pow'].to_numpy()

N = data['N'].to_numpy()
N_err = data['N_err'].to_numpy()


N = N - N[-1]

# Fit the data
p0 = [0, -10 / 150, 150]

x = bar_pow[bar_pow >= 1]
y = N[bar_pow >= 1]
y_fit = y[x <= 3500]
x_fit = x[x <= 3500]

popt, pcov = curve_fit(lin, x_fit, y_fit, p0=p0)

pl.hlines(0, 100, 18750, label='Zero line', ls='--', color=RPTU_COLORS["schiefer"])

def nullstelle(x0,m,off):
    return -off/m +x0

x_0 = nullstelle(*popt)

print('Zero Line crossing y=0 ',x_0 )
b = np.arange(0,5000,1)

# plot fit
pl.plot(b,lin(b,*popt),color = RPTU_COLORS["himbeere"],label=fr'Fit' ) # $x_0$ = {x_0 :.2f} ')

# plot data
pl.errorbar(bar_pow, N, yerr=N_err, **get_style(RPTU_COLORS["pflaume"]), label = 'Data')

pl.ylabel(r'Remaining atoms N')
pl.xlabel(r'Barrier Intensity  $I$ [arb. units]')

pl.legend()
pl.ylim(-100,2000)
pl.tight_layout()

pl.show()