# coding: utf-8
from __future__ import print_function, division, unicode_literals

import numpy               as np
import uncertainties.umath as um
import matplotlib.pyplot   as plt
from lab import *

##
## Part I (charge)
##

##
## overdamped oscillation
##

R  = ufloat(9.93e2,  0.05e2) # resistance
L  = ufloat(3.82e-3, 0.5e-3) # inductance
C  = ufloat(8.32e-6, 0.5e-6) # capacitance
V0 = ufloat(10, 0)           # generator peak tension
nu = ufloat(2, 0)            # generator frequency

# time
t = array(0.2, 0.4, 0.6, 0.8, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,
          12, 13, 14, 15, 16, 17, 18, 19, 20)*1e-3

# tension
V = array(9.8, 9.56, 9.32, 9.09, 8.89, 7.87, 6.97, 6.17, 5.47,
          4.85, 4.3, 3.81, 3.37, 2.99, 2.65, 2.34, 2.09, 1.83,
          1.64, 1.46, 1.27, 1.14, 1.01, 0.9)

# fit V(t)
w0 = 1/um.sqrt(L*C)
y  = R/(2*L)
a  = V0
b  = um.sqrt(y**2 - w0**2)

def over(t, a, b, y):
  return a * y/b * (np.exp(-(y-b)*t) - np.exp(-(y+b)*t))

x = np.linspace(t.min(), t.max(), 500)
f, p = curve(t, V, over, 0.01, guess=[V0.n, b.n, y.n])

# χ² and compatibility tests
alpha = chi_squared_fit(t, V, f, sigma=0.01, s=3)
beta  = [check_measures(*i) for i in zip([a, b, y], p)]

print(mformat('''
overdamped oscillation fit
α: {}V,  β₀={:.2f}
β: {}Hz, β₁={:.2f}
γ: {}Hz, β₂={:.2f}

χ²: α={:.2f}, α>ε: {}
''', p[0], beta[0]
   , p[1], beta[1]
   , p[2], beta[2]
   , alpha, alpha>epsilon))

# plot V(t)
plt.figure(1)
plt.title('overdamped oscillations (charge)')
plt.xlabel('time (ms)')
plt.ylabel('tension (V)')
plt.scatter(t*1e3, V, color='#5c6652')
plt.plot(x*1e3, f(x), color='#718062')
plt.show()


##
## underdamped oscillations
##

R  = ufloat(100, 4)           # resistance
L  = ufloat(3.82e-3, 0.5e-03) # inductance
C  = ufloat(44.5e-9, 0.5e-9)  # capacitance
V0 = ufloat(5, 0)             # generator peak tension
nu = ufloat(1, 0)             # generator frequency

# time
t = np.append(np.arange(0, 100, 2), np.arange(104, 250, 4))*1e-6

# tension
V = array(0.285,  0.522,  0.712,  0.879,  1.040,  1.158,  1.229,  1.286,
          1.292,  1.256,  1.214,  1.117,  1.026,  0.900,  0.777,  0.605,
          0.437,  0.271,  0.121, -0.045, -0.196, -0.335, -0.453, -0.564,
         -0.636, -0.701, -0.758, -0.777, -0.770, -0.766, -0.715, -0.673,
         -0.602, -0.524, -0.450, -0.351, -0.258, -0.142, -0.062,  0.041,
          0.122,  0.211,  0.268,  0.341,  0.382,  0.426,  0.454,  0.466,
          0.471,  0.453,  0.366,  0.259,  0.135,  0.017, -0.081, -0.165,
         -0.233, -0.270, -0.287, -0.249, -0.217, -0.164, -0.078, -0.003,
          0.052,  0.109,  0.146,  0.147,  0.159,  0.146,  0.122,  0.100,
          0.050, -0.002, -0.034, -0.070, -0.073, -0.100, -0.115, -0.093,
         -0.061, -0.054, -0.023, -0.014,  0.023,  0.051,  0.059)


# fit V(t)
a  = R*C*V0
y  = R/(2*L)
w0 = 1/um.sqrt(L*C)
w  = um.sqrt(w0**2 - y**2)

def under(t, a, w, y):
  return a * np.exp(-y*t) * (y*np.cos(w*t) + w*np.sin(w*t))

x = np.linspace(0, t.max(), 500)
f, p = curve(t, V, under, 0.01, guess=[a.n, w.n, y.n])

# χ² and compatibility tests
alpha = chi_squared_fit(t, V, f, sigma=0.01, s=3)
beta  = [check_measures(*i) for i in zip([a, w, y], p)]

print(mformat('''
underdamped oscillation fit
α: {}Wb, β₀={:.2f}
ω: {}Hz, β₁={:.2f}
γ: {}Hz, β₂={:.2f}

χ²: α={:.2f}, α>ε: {}
''', p[0], beta[0]
   , p[1], beta[1]
   , p[2], beta[2]
   , alpha, alpha>epsilon))

plt.figure(2)
plt.title('underdamped oscillations (charge)')
plt.xlabel('time (μs)')
plt.ylabel('tension (V)')
plt.scatter(t*1e6, V, s=12, color='#a17e3e')
plt.plot(x*1e6, f(x), color='#c8b491')
plt.show()


##
## critical dampening
##

R  = ufloat(586, 4)           # resistance
L  = ufloat(3.82e-3, 0.5e-3)  # inductance
C  = ufloat(4.45e-8, 0.05e-8) # capacitance
V0 = ufloat(2.5, 0)           # generator peak tension

# time
t = array(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 14, 18, 22, 26, 30, 34,
          38, 42, 46, 50, 54, 58, 62, 66, 70, 74, 78, 82, 86, 90,
          94, 98, 102, 106, 110, 114, 118, 122, 126, 130, 134, 138, 
          142, 146, 150)*1e-6
# tension
V = array(0.007,  0.355,  0.654,  0.907,  1.120,  1.300,  1.437,  1.560,
          1.639,  1.726,  1.772,  1.821,  1.726,  1.557,  1.355,  1.145,
          0.954,  0.783,  0.642,  0.526,  0.411,  0.328,  0.254,  0.213,
          0.157,  0.125,  0.095,  0.069,  0.057,  0.040,  0.034,  0.026,
          0.025,  0.024,  0.013,  0.014,  0.003,  0.010,  0.010,  0.007,
          -0.001, 0.002, -0.004, -0.006, -0.008,  0.002)


## fit as critical

y = R/(2*L)
a = R*C*V0

def critic(t, a, y):
  return a * y**2 * t*np.exp(-y*t)

x = np.linspace(0, t.max(), 500)
f, p = curve(t, V, critic, 0.005, guess=[a.n, y.n])

# χ² and compatibility tests
alpha = chi_squared_fit(t, V, f, sigma=0.005, s=2)
beta  = [check_measures(*i) for i in zip([a, y], p)]

print(mformat('''
critically damped oscillation fit
α: {}Wb, β₀={:.2f}
γ: {}Hz, β₂={:.2f}

χ²: α={:.2f}, α>ε: {}
''', p[0], beta[0]
   , p[1], beta[1]
   , alpha, alpha>epsilon))

# plot V(t)
plt.figure(3)
plt.subplot(3, 1, 1)
plt.title('critically damped oscillations (charge)')
plt.xlabel('time (μs)')
plt.ylabel('tension (V)')
plt.scatter(t*1e6, V, color='#43454f')
plt.plot(x*1e6, f(x), color='#65788f')


## fit as overdamped

w0 = 1/um.sqrt(L*C)
y  = R/(2*L)
a  = V0
b  = um.sqrt(y**2 - w0**2)

def over(t, a, b, y):
  return a * y/b * (np.exp(-(y-b)*t) - np.exp(-(y+b)*t))

x = np.linspace(t.min(), t.max(), 500)
f, p = curve(t, V, over, 0.01, guess=[V0.n, b.n, y.n])

# χ² and compatibility tests
alpha = chi_squared_fit(t, V, f, sigma=0.01, s=3)
beta  = [check_measures(*i) for i in zip([a, b, y], p)]

print(mformat('''
critically damped oscillation fit (as overdamped)
α: {}V,  β₀={:.2f}
β: {}Hz, β₁={:.2f}
γ: {}Hz, β₂={:.2f}

χ²: α={:.2f}, α>ε: {}
''', p[0], beta[0]
   , p[1], beta[1]
   , p[2], beta[2]
   , alpha, alpha>epsilon))

# plot V(t)
plt.subplot(3, 1, 2)
plt.xlabel('time (ms)')
plt.ylabel('tension (V)')
plt.scatter(t*1e3, V, color='#5c6652')
plt.plot(x*1e3, f(x), color='#718062')


## fit as underdamped

a  = R*C*V0
y  = R/(2*L)
w0 = 1/um.sqrt(L*C)
w  = um.sqrt(abs(w0**2 - y**2))

def under(t, a, w, y):
  return a * np.exp(-y*t) * (y*np.cos(w*t) + w*np.sin(w*t))

x = np.linspace(0, t.max(), 500)
f, p = curve(t, V, under, 0.01, guess=[a.n, w.n, y.n])

# χ² and compatibility tests
alpha = chi_squared_fit(t, V, f, sigma=0.01, s=3)
beta  = [check_measures(*i) for i in zip([a, w, y], p)]

print(mformat('''
critically damped oscillation fit (as underdamped)
α: {}Wb, β₀={:.2f}
ω: {}Hz, β₁={:.2f}
γ: {}Hz, β₂={:.2f}

χ²: α={:.2f}, α>ε: {}
''', p[0], beta[0]
   , p[1], beta[1]
   , p[2], beta[2]
   , alpha, alpha>epsilon))

plt.subplot(3,1,3)
plt.xlabel('time (μs)')
plt.ylabel('tension (V)')
plt.scatter(t*1e6, V, s=12, color='#a17e3e')
plt.plot(x*1e6, f(x), color='#c8b491')
plt.show()