# %%
import matplotlib.pyplot as plt
import RigolWFM.wfm as rigol
import numpy as np
import matplotlib as mpl
import scipy.optimize

# %%
# Data import

f = np.array(
    [1, 1.83, 3.36, 6.16, 11.3, 20.7, 37.9, 69.5, 127, 234, 428, 785, 1440, 2640, 4830, 8860, 16200, 29800, 54600])

A_HH = np.array(
    [1.13, 1.13, 1.13, 1.13, 1.13, 1.14, 1.13, 1.11, 1.03, 0.881, 0.708, 0.504, 0.334, 0.198, 0.119, 0.069, 0.046,
     0.028, 0.024])
dA_HH = np.array(
    [0.006, 0.002, 0.001, 0.002, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.001, 0.002, 0.001, 0.004, 0.001, 0.001,
     0.001, 0.001, 0.001])

ph_HH = np.array(
    [0.8, -0.8, -1.2, -0.8, 1.5, 3.7, 7.9, 11.8, 20.6, 32.0, 44.6, 54.8, 64.3, 74.0, 78.6, 82.8, 85.3, 87.4, 93.7])
dph_HH = np.array(
    [1.5, 1.5, 0.5, 1.2, 4.1, 0.001, 1.3, 0.001, 0.4, 0.001, 0.080, 0, 0.23, 0.7, 0.05, 1.6, 1.5, 2.5, 2.8])

A_AHH = np.array(
    [1.14, 1.14, 1.14, 1.14, 1.15, 1.15, 1.13, 1.09, 0.966, 0.737, 0.539, 0.359, 0.216, 0.130, 0.0767, 0.0475])
dA_AHH = np.array(
    [0.007, 0.005, 0.002, 0.002, 0.002, 0.002, 0.001, 0.001, 0.004, 0.012, 0.001, 0.001, 0.003, 0.001, 0.0004, 0.0004])

ph_AHH = np.array([0.6, -0.3, 0.5, 0.3, 2.5, 4.3, 8.8, 17.2, 30.3, 45.3, 54.2, 64.4, 72.6, 77.24, 79.4, 80.14])
dph_AHH = np.array([3.0, 2.0, 0.4, 1.2, 1.0, 0.7, 1.2, 1.2, 0.5, 2.6, 0.9, 1.4, 1.2, 1.5, 2.1, 1.2])

# %%
my_colors = {'light_green': '#97e144',
             'orange': '#FF914D',
             'light_grey': '#545454',
             'pastel_blue': '#1b64d1',
             'light_blue': '#71C8F4',
             'purple': '#7c588c'}


mpl.rcParams.update({'font.size': 11, 'axes.linewidth': 1, 'lines.linewidth': 2, 'text.usetex': False, 'font.family': 'arial'})
mpl.rcParams['xtick.direction'] = 'in'
mpl.rcParams['ytick.direction'] = 'in'

mpl.rcParams['xtick.top'] = True
mpl.rcParams['ytick.right'] = True
# %% Fit func

def I_fit(f,U,tau):
    omega = 2* np.pi *f
    return U/np.sqrt(tau**2 + omega**2)
# %% fit HH

fit_bound = (0,len(A_HH))
p0 = [1,2.2e3]
popt_hh, pcov_hh = scipy.optimize.curve_fit(I_fit, f, A_HH, p0=p0, sigma=dA_HH, absolute_sigma=False)
print("U0, tau, t0")
print(popt_hh)

# %% fit AHH

fit_bound = (0,len(A_HH))
p0 = [1,2.2e3]
popt_ahh, pcov_ahh = scipy.optimize.curve_fit(I_fit, f[0:len(A_AHH)], A_AHH, p0=p0)
print("U0, tau, t0")
print(popt_hh)
# %% Scaling factor
scale = 0.3/I_fit(1, *popt_hh)

# %%
x=np.linspace(0,50000,10000)
fig, ax = plt.subplots(figsize=(5.4, 4), dpi=400)

ax.errorbar(f, scale*A_HH, yerr=scale*dA_HH, elinewidth=2, capsize=5, linewidth=0, label='Data HH', zorder=4, marker='.',color='C0')
ax.plot(x, scale*I_fit(x, *popt_hh),label=f'fit HH, $\\tau$ = 2255(110) 1/s',color='C8')
ax.errorbar(f[0:len(A_AHH)], scale*A_AHH, yerr=scale*dA_AHH, elinewidth=2, capsize=5, linewidth=0, label='Data AHH', zorder=3, marker='^',color='C1')
ax.plot(x, scale*I_fit(x, *popt_ahh),label=f'fit AHH, $\\tau$ = 1416(56) 1/s', color='C4')

ax.hlines(scale*I_fit(1, *popt_hh)/np.sqrt(2),0, 1e5, color='C7', linestyle=(0,(2.5,3)), label='3dB point \nHH: f = 359 Hz, AHH: 225 Hz ')
ax.grid(alpha=0.6)
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlim(0, 1e5)
ax.set_ylabel(r'current amplitude ($\rm A_{pp}$)')
ax.set_xlabel('frequency f (Hz)')
handles, labels = ax.get_legend_handles_labels()
ax.legend(handles=[handles[3], handles[4], handles[0], handles[1], handles[2]])

fig.tight_layout()
fig.savefig('C:/Users/Joschka/Desktop/Master_Thesis/Figures/Coil_measurements/Final/freq_resp.png')
fig.savefig('C:/Users/Joschka/Desktop/Master_Thesis/Figures/Coil_measurements/Final_low/freq_resp.png', dpi=96)
plt.show()
# %% Find 3dB point
x=np.linspace(0,1000,10000)
for i in range(0,10000):
    if I_fit(x[i], *popt_hh) < (I_fit(1, *popt_hh)/np.sqrt(2)):
        print(f"Cutoff HH: {x[i]} Hz")
        break

for i in range(0, 10000):
    if I_fit(x[i], *popt_ahh) < (I_fit(1, *popt_ahh) / np.sqrt(2)):
        print(f"Cutoff AHH: {x[i]} Hz")
        break
# %%
fig, ax = plt.subplots(figsize=(5.4, 4),dpi=400)

ax.errorbar(f, ph_HH, yerr=dph_HH, label='Data HH')


ax.grid(alpha=0.6, which='minor')
ax.set_xscale('log')
#ax.set_yscale('log')

plt.show()