import numpy as np
import matplotlib.pyplot as pl
import pandas as pd
import scipy.constants as c

from style import  *
from helper import *

import os
from scipy.optimize import curve_fit

def func(N, A):
    return A * N ** (1 / 2)

path_data_repo = "../data_supplementary"

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

path = os.path.join(path_data_repo ,"chem_pot_atomnumber_GPE.csv")

data = pd.read_csv(path)

N = data["atomnumber"]
mu = data["mu"] / c.h

p0 = [1]
popt, pcov = curve_fit(func, N, mu, p0=p0)

n = np.arange(0, 200e3, 1)

pl.errorbar(n*1e-3, func(n, *popt), color=RPTU_COLORS['mango'] , label=fr'Fit: $f(N) = {popt[0]:.3f} \mathrm{{Hz}} \cdot N ^{{1/2}}$',zorder = 2)

pl.errorbar(N * 1e-3, mu, **get_style(RPTU_COLORS['nacht']),zorder = 1)
pl.xlabel(r'Atomnumber $N$ [$\times 10³ $]')
pl.ylabel(r'Chemical Potential $µ$ [Hz]')
pl.legend()
pl.tight_layout()
#pl.savefig(f'fit_mu.png')
pl.show()