import numpy as np
import matplotlib.pyplot as plt
import scipy.odr as sodr
import scipy.stats as scs
from my_methods import *

delta_U = 0.00027#20/(2**12) /np.sqrt(12)
errorscale = 30
R = 996 #Ohm
sigma_R = 1/np.sqrt(12) #Ohm

kondensator_aufladung = parseCSV('kondensator_aufladung')
kondensator_entladung = parseCSV('kondensator_entladung')

t_auf = kondensator_aufladung[:, 0]
U_auf = kondensator_aufladung[:, 1]
t_ent = kondensator_entladung[:, 0]
U_ent = kondensator_entladung[:, 1]

sigma_U_auf = delta_U * np.ones_like(U_auf)
sigma_U_ent = delta_U * np.ones_like(U_ent)
sigma_t_auf = (t_auf[1] - t_auf[0]) / np.sqrt(12) * np.ones_like(U_auf)
sigma_t_ent = (t_ent[1] - t_ent[0]) / np.sqrt(12) * np.ones_like(U_ent)

plt.errorbar(t_auf, U_auf, xerr = sigma_t_auf, yerr = sigma_U_auf * errorscale, label = 'Aufladung', linestyle = '')
plt.errorbar(t_ent, U_ent, xerr = sigma_t_ent, yerr = sigma_U_ent * errorscale, label = 'Entladung', linestyle = '')
plt.ylabel(r'$U_A [V]$')
plt.xlabel(r'$t [ms]$')
plt.title(r'Messdaten Kondensator$')
plt.grid()
plt.legend()
plt.tight_layout()
plt.savefig('output/Kondensator_messung.pdf')
plt.clf()

print('########## Aufladung ##########')
output = exponential_fit(t_auf, U_auf, sigma_t_auf, sigma_U_auf, Lambda0 = -1e-3, A0 = -5, offset0 = 5)
chiq_dof, p_value = chi_sq_dof_p(t_auf, U_auf, sigma_t_auf, sigma_U_auf, output.beta, exponential)
print(output.beta)
print(np.sqrt(np.diagonal(output.cov_beta)))

Lambda = -output.beta[0] * 1e3
sigma_Lambda = np.sqrt(output.cov_beta[0, 0]) * 1e3
C = 1 / (Lambda * R)
sigma_C = ((sigma_R/(Lambda * R**2))**2 + (sigma_Lambda/(Lambda**2 * R))**2)**0.5
print(f'\n\t Lambda: {round(Lambda, sigma_Lambda)}')
print(f'\t C: {round(C * 1e6, sigma_C * 1e6)} µF\n')

fig, ax = plt.subplots(nrows = 2, sharex = True)
ax[0].plot(t_auf, exponential(output.beta, t_auf), label = r'fit $y = K + A \cdot e^{\Lambda \cdot t}$', color='C1')
ax[0].errorbar(t_auf[::100], U_auf[::100], xerr = sigma_t_auf[::100], yerr = sigma_U_auf[::100], label='data', linestyle = '', marker = ',')
ax[1].errorbar(t_auf, U_auf - exponential(output.beta, t_auf), xerr = sigma_t_auf, yerr = sigma_U_auf, label = 'residuals', linestyle = '', marker = ',')
ax[0].set(ylabel = r'$U_A [V]$', title = r'Fit für Aufladung des Kondensators')
ax[1].set(xlabel = r'$t [ms]$', ylabel = r'Data - Fit [V]')
ax[0].text(0.982, 0.38, r'$\lambda$'+f': {round(output.beta[0], np.sqrt(output.cov_beta[0, 0]), scientific = True)}\n'+r'$A$'+f': {round(output.beta[1], np.sqrt(output.cov_beta[1, 1]))}\n'+r'$K$'+f': {round(output.beta[2], np.sqrt(output.cov_beta[2, 2]))}\n'+r'$\frac{\chi^2}{dof}$'+f': {np.round(chiq_dof, 1)}\np: {p_value}' , horizontalalignment='right', verticalalignment='bottom', transform=ax[0].transAxes, bbox=dict(facecolor='white', edgecolor='gray'))
ax[0].grid()
ax[1].grid()
ax[0].legend(loc = 'lower right')
ax[1].legend()
fig.tight_layout()
fig.savefig(f'./output/kondensator_aufladung_fit.pdf')
plt.clf()

print('########## Entladung ##########')
output = exponential_fit(t_ent, U_ent, sigma_t_ent, sigma_U_ent, Lambda0 = -1e-3, A0 = 5, offset0 = -5)
chiq_dof, p_value = chi_sq_dof_p(t_ent, U_ent, sigma_t_ent, sigma_U_ent, output.beta, exponential)
print(output.beta)
print(np.sqrt(np.diagonal(output.cov_beta)))

Lambda = -output.beta[0] * 1e3
sigma_Lambda = np.sqrt(output.cov_beta[0, 0]) * 1e3
C = 1 / (Lambda * R)
sigma_C = ((sigma_R/(Lambda * R**2))**2 + (sigma_Lambda/(Lambda**2 * R))**2)**0.5
print(f'\n\t Lambda: {round(Lambda, sigma_Lambda)}')
print(f'\t C: {round(C * 1e6, sigma_C * 1e6)} µF\n')

fig, ax = plt.subplots(nrows = 2, sharex = True)
ax[0].plot(t_ent, exponential(output.beta, t_ent), label = r'fit $y = K + A \cdot e^{\Lambda \cdot t}$', color='C1')
ax[0].errorbar(t_ent[::100], U_ent[::100], xerr = sigma_t_ent[::100], yerr = sigma_U_ent[::100], label='data', linestyle = '', marker = ',')
ax[1].errorbar(t_ent, U_ent - exponential(output.beta, t_ent), xerr = sigma_t_ent, yerr = sigma_U_ent, label = 'residuals', linestyle = '', marker = ',')
ax[0].set(ylabel = r'$U_A [V]$', title = r'Fit für Entladung des Kondensators')
ax[1].set(xlabel = r'$t [ms]$', ylabel = r'Data - Fit [V]')
ax[0].text(0.982, 0.12, r'$\lambda$'+f': {round(output.beta[0], np.sqrt(output.cov_beta[0, 0]), scientific = True)}\n'+r'$A$'+f': {round(output.beta[1], np.sqrt(output.cov_beta[1, 1]))}\n'+r'$K$'+f': {round(output.beta[2], np.sqrt(output.cov_beta[2, 2]))}\n'+r'$\frac{\chi^2}{dof}$'+f': {np.round(chiq_dof, 1)}\np: {p_value}' , horizontalalignment='right', verticalalignment='bottom', transform=ax[0].transAxes, bbox=dict(facecolor='white', edgecolor='gray'))
ax[0].grid()
ax[1].grid()
ax[0].legend(loc = 'upper right')
ax[1].legend()
fig.tight_layout()
fig.savefig(f'./output/kondensator_entladung_fit.pdf')
plt.clf()