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

delta_U = 0.00027#20/(2**12) /np.sqrt(12)
errorscale = 1

print('\n########## Verstaerker ohne C3 ##########')
verstaerker = parseCSV('verstaerker_ohne_c3')
t = verstaerker[:, 0]
U_a = verstaerker[:, 1]
U_b = verstaerker[:, 2]

sigma_t = (t[1] - t[0])/np.sqrt(12) * np.ones_like(t)
sigma_U_a = delta_U * np.ones_like(U_a)
sigma_U_b = delta_U * np.ones_like(U_b)

errorscale = 50
plt.errorbar(t, U_a, xerr = sigma_t, yerr = sigma_U_a * errorscale, label = 'U_A', linestyle = '')
plt.errorbar(t, U_b, xerr = sigma_t, yerr = sigma_U_b * errorscale, label = 'U_B', linestyle = '')
plt.ylabel(r'$U [V]$')
plt.xlabel(r'$t [ms]$')
plt.title(r'Messdaten Gleichrichter')
plt.grid()
plt.legend()
plt.tight_layout()
plt.savefig('output/verstaerker_oszi_ohne_c3.pdf')
plt.clf()

print('########## fitting sine ##########')
# odr does not converge here :(
def sine2(x, a,b,c,d):
    return sine([a,b,c,d], x)

popt_a, pcov_a = scipy.optimize.curve_fit(sine2, t, U_a, sigma = sigma_U_a, p0 = [2*np.pi, np.pi/2, 0.5, 0])
popt_b, pcov_b = scipy.optimize.curve_fit(sine2, t, U_b, sigma = sigma_U_b, p0 = [2*np.pi, -np.pi/2, 2, 0])
chi_a, p_a = chi_sq_dof_p(t, U_a, xerr = sigma_t, yerr = sigma_U_a, beta = popt_a, model = sine)
chi_b, p_b = chi_sq_dof_p(t, U_b, xerr = sigma_t, yerr = sigma_U_b, beta = popt_b, model = sine)
sigma_a = np.sqrt(pcov_a.diagonal())
sigma_b = np.sqrt(pcov_b.diagonal())

print('## U_a ##')
print(popt_a)
print(sigma_a)
print(chi_a, p_a)
print('## U_b ##')
print(popt_b)
print(sigma_b)
print(chi_b, p_b)

beta = popt_b[2] / popt_a[2]
sigma_beta = np.sqrt((beta*sigma_b[2]/popt_b[2])**2 + (beta*sigma_a[2]/popt_a[2])**2)

print(f'\n\t Verstärkung: {round(beta, sigma_beta)}')

errorscale = 1
plt.errorbar(t, U_a, xerr = sigma_t, yerr = sigma_U_a * errorscale, label = 'U_A', linestyle = '')
plt.errorbar(t, U_b, xerr = sigma_t, yerr = sigma_U_b * errorscale, label = 'U_B', linestyle = '')
plt.plot(t, sine(popt_a, t), label = 'a')
plt.plot(t, sine(popt_b, t), label = 'b')
plt.legend()
plt.grid()
plt.savefig('./output/verstaerker_ohne_c3_fit.pdf')
plt.cla()

print('\n\n########## Verstaerker mit C3 ##########')
verstaerker = parseCSV('verstaerker_mit_c3')
t = verstaerker[:, 0]
U_a = verstaerker[:, 1]
U_b = verstaerker[:, 2]

sigma_t = (t[1] - t[0])/np.sqrt(12) * np.ones_like(t)
sigma_U_a = delta_U * np.ones_like(U_a)
sigma_U_b = delta_U * np.ones_like(U_b)

errorscale = 50
plt.errorbar(t, U_a, xerr = sigma_t, yerr = sigma_U_a * errorscale, label = 'U_A', linestyle = '')
plt.errorbar(t, U_b, xerr = sigma_t, yerr = sigma_U_b * errorscale, label = 'U_B', linestyle = '')
plt.ylabel(r'$U [V]$')
plt.xlabel(r'$t [ms]$')
plt.title(r'Messdaten Gleichrichter')
plt.grid()
plt.legend()
plt.tight_layout()
plt.savefig('output/verstaerker_oszi_ohne_c3.pdf')
plt.clf()

print('########## fitting sine ##########')
# odr does not converge here :(
def sine2(x, a,b,c,d):
    return sine([a,b,c,d], x)

popt_a, pcov_a = scipy.optimize.curve_fit(sine2, t, U_a, sigma = sigma_U_a, p0 = [2*np.pi, np.pi/2, 0.05, 0])
popt_b, pcov_b = scipy.optimize.curve_fit(sine2, t, U_b, sigma = sigma_U_b, p0 = [2*np.pi, -np.pi/2, 3.55, -0.3])
chi_a, p_a = chi_sq_dof_p(t, U_a, xerr = sigma_t, yerr = sigma_U_a, beta = popt_a, model = sine)
chi_b, p_b = chi_sq_dof_p(t, U_b, xerr = sigma_t, yerr = sigma_U_b, beta = popt_b, model = sine)
sigma_a = np.sqrt(pcov_a.diagonal())
sigma_b = np.sqrt(pcov_b.diagonal())

print('## U_a ##')
print(popt_a)
print(sigma_a)
print(chi_a, p_a)
print('## U_b ##')
print(popt_b)
print(sigma_b)
print(chi_b, p_b)

beta = popt_b[2] / popt_a[2]
sigma_beta = np.sqrt((beta*sigma_b[2]/popt_b[2])**2 + (beta*sigma_a[2]/popt_a[2])**2)

print(f'\n\t Verstärkung: {round(beta, sigma_beta)}')

errorscale = 1
plt.errorbar(t, U_a, xerr = sigma_t, yerr = sigma_U_a * errorscale, label = 'U_A', linestyle = '')
plt.errorbar(t, U_b, xerr = sigma_t, yerr = sigma_U_b * errorscale, label = 'U_B', linestyle = '')
plt.plot(t, sine(popt_a, t), label = 'a')
plt.plot(t, sine(popt_b, t), label = 'b')
plt.legend()
plt.grid()
plt.savefig('./output/verstaerker_mit_c3_fit.pdf')
plt.show()
plt.cla()