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

import matplotlib as mpl
# %%
import scipy.optimize

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

# %%
hhpion.channels[1].volts

# %%

scope = 'DS1104Z-S'

hhpion = rigol.Wfm.from_file('C:/Users/Joschka/Desktop/Coil_Data/New/hhpion.wfm', scope)
hhpioff = rigol.Wfm.from_file('C:/Users/Joschka/Desktop/Coil_Data/New/hhpioff.wfm', scope)
hhon = rigol.Wfm.from_file('C:/Users/Joschka/Desktop/Coil_Data/New/hhon.wfm', scope)
hhoff = rigol.Wfm.from_file('C:/Users/Joschka/Desktop/Coil_Data/New/hhoff.wfm', scope)

ahhpion = rigol.Wfm.from_file('C:/Users/Joschka/Desktop/Coil_Data/New/ahhpion.wfm', scope)
ahhpioff = rigol.Wfm.from_file('C:/Users/Joschka/Desktop/Coil_Data/New/ahhpioff.wfm', scope)
ahhoff = rigol.Wfm.from_file('C:/Users/Joschka/Desktop/Coil_Data/New/ahhoff.wfm', scope)
ahhon = rigol.Wfm.from_file('C:/Users/Joschka/Desktop/Coil_Data/New/ahhon.wfm', scope)

# %%
print(hhon.channels[1])

# %%
def I(V):
    return (V/1.2093 -2.47e-3) * 1.21/1.276

print(I(1.21))
print(1/0.88)


def I_exp(t, I_0, tau, t0):
    return I_0 *(1-np.exp(-tau *(t-t0)))

def I_pi(t, tau, t0=0):
    return 30/3.2 * (1 - np.exp(-tau * (t - t0)))
# %%
x = np.linspace(0,80e-6,1000)
fig, ax = plt.subplots(figsize=(5, 3.6), dpi=400)


ax.plot(1e6* hhpion.channels[1].times, I(hhpion.channels[1].volts), label='HH on')
ax.plot(1e6 *ahhpion.channels[1].times, I(ahhpion.channels[1].volts),label='AHH on')

# ax.plot(1e6 *x, I_pi(x, popt_hhon[1]))


ax.plot(1e6* hhpioff.channels[1].times, I(hhpioff.channels[1].volts), label='HH off')
ax.plot(1e6 *ahhpioff.channels[1].times, I(ahhpioff.channels[1].volts), label='AHH off')
ax.grid(alpha=0.5)
ax.set_xlabel('time (μs)')
ax.set_ylabel('current I (A)')

ax.set_xlim(-50, 300)
ax.legend()
fig.tight_layout()
plt.savefig('C:/Users/Joschka/Desktop/Master_Thesis/Figures/Coil_measurements/Final/time_resp_pi.png', dpi=400)
plt.savefig('C:/Users/Joschka/Desktop/Master_Thesis/Figures/Coil_measurements/Final_low/time_resp_pi.png',dpi=96)
plt.show()
# %%
# Fitting
# 0 at 745128
fitx = np.array(hhon.channels[1].times)
fity = np.array(I(hhon.channels[1].volts))

dfity = np.array(hhon.channels[1].volts)
dfity = 0.1*200e-3 + 2e-3 + 0.001 * dfity
dfity = I(dfity)
fit_bound = (745128,len(fitx))
p0 = [1,2.2e3,0]
popt_hhon, pcov_hhon = scipy.optimize.curve_fit(I_exp, fitx[fit_bound[0]:fit_bound[1]], fity[fit_bound[0]:fit_bound[1]],
                                                p0=p0, sigma = dfity[fit_bound[0]:fit_bound[1]], absolute_sigma=True)
print("I0, tau, t0")
print(popt_hhon)
# %%
# Fitting
# 0 at 510128
fitx = np.array(ahhon.channels[1].times)
fity = np.array(I(ahhon.channels[1].volts))
fit_bound = (510128,len(fitx))
p0 = [1,2.2e3,0]
popt_ahhon, pcov_ahhon = scipy.optimize.curve_fit(I_exp, fitx[fit_bound[0]:fit_bound[1]], fity[fit_bound[0]:fit_bound[1]], p0=p0)
print("I0, tau, t0")
print(popt_ahhon)
# %%
for i in range(0,len(fitx)):
    if fitx[i]>0:
        print(i)
        break
# %%
x = np.linspace(0, 0.01, 1000)
fig, ax = plt.subplots(figsize=(5, 3.6))


ax.scatter(1e3* fitx,fity, label='HH on', linewidth=0.01)

ax.legend()
plt.show()
# %%
x = np.linspace(0, 0.01, 1000)
fig, ax = plt.subplots(figsize=(5, 3.6),dpi=400)


ax.plot(1e3* hhon.channels[1].times, I(hhon.channels[1].volts), label='HH on', color='C0')
ax.plot(1e3 *ahhon.channels[1].times, I(ahhon.channels[1].volts),label='AHH on', color='C1')

ax.plot(1e3*x, I_exp(x, *popt_hhon))

ax.plot(1e3* hhoff.channels[1].times, I(hhoff.channels[1].volts), label='HH off', color='C2',zorder=0)
ax.plot(1e3 *ahhoff.channels[1].times, I(ahhoff.channels[1].volts), label='AHH off', color='C3',zorder=1)

ax.plot(1e3*x, I_exp(x, *popt_hhon), color = 'C8', linestyle=(0,(2,)),
        label=f'fit HH on, $ \\tau$ = {popt_hhon[1]:.0f} 1/s', zorder=3)
ax.plot(1e3*x, I_exp(x, *popt_ahhon), color = 'C4', linestyle=(0,(2,)),
        label=f'fit AHH on, $\\tau$ = {popt_ahhon[1]:.0f} 1/s')
ax.grid(alpha=0.5)
ax.set_xlabel('time (ms)')
ax.set_ylabel('current I (A)')

ax.legend(loc=5)
fig.tight_layout()
plt.savefig('C:/Users/Joschka/Desktop/Master_Thesis/Figures/Coil_measurements/Final/time_resp.png')
plt.savefig('C:/Users/Joschka/Desktop/Master_Thesis/Figures/Coil_measurements/Final_low/time_resp.png', dpi=96)
plt.show()


# %%
# Resistance HH Calculation
I = 1
U = 3.32
dI = 1e-3 + 0.001*I
dU = np.sqrt(5e-3**2 + (0.0005*U)**2 + (0.005*U)**2)
print(dU)
#dU = 0.05

R =U/I
dR = R* np.sqrt((dI/I)**2 + (dU/U)**2)
print(f"R = {R:.3f} +/- {dR:.3f} Ω")

L = R/popt_hhon[1]
dL=dR/popt_hhon[1]
print(f"L = {L:.6f} +/-{dL:.6f} Ω")

# %%
# Resistance AHH Calculation
I = 1
U = 6.64
dI = np.sqrt(1e-3**2 + (0.001*I)**2)
dU = np.sqrt(5e-3**2 + (0.0005*U)**2 + (0.005*U)**2)
print(dU)

R =U/I
dR = R* np.sqrt((dI/I)**2 + (dU/U)**2)
print(f"R = {R:.3f} +/- {dR:.3f} Ω")

L = R/popt_ahhon[1]
dL=dR/popt_ahhon[1]
print(f"L = {L:.5f} +/-{dL:.5f} Ω")

# %%