# 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 (discharge)
##

##
## 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(5, 0)            # generator peak tension
nu = ufloat(2, 0)            # generator frequency

# time
t = array(0,   40e-6,   80e-6,  120e-6,  160e-6,  200e-6,  240e-6,  280e-6,  320e-6,
     360e-6,  400e-6,  440e-6,  480e-6,  520e-6,  560e-6,  600e-6,  640e-6,  680e-6,
     720e-6,  760e-6,  800e-6,  840e-6,  880e-6,  920e-6,  960e-6,    1e-3, 1.04e-3,
    1.08e-3, 1.12e-3, 1.16e-3, 1.24e-3, 1.32e-3, 1.48e-3, 1.64e-3, 1.88e-3, 2.12e-3,
    2.32e-3, 2.68e-3, 3.00e-3, 3.56e-3)

# tension
V = array(0.0, -9.0, -8.6, -8.2, -7.9, -7.6, -7.2, -6.9, -6.6, -6.3, -6.0, -5.7, -5.5,
         -5.2, -5.0, -4.8, -4.6, -4.4, -4.2, -4.0, -3.8, -3.6, -3.5, -3.3, -3.2, -3.0,
         -2.9, -2.8, -2.7, -2.5, -2.3, -2.1, -1.8, -1.5, -1.1, -0.8, -0.7, -0.5, -0.3,
         -0.18)

# 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 (discharge)')
plt.xlabel('time (ms)')
plt.ylabel('tension (V)')
plt.scatter(t*1e3, V, color='#8c4f4a')
plt.plot(x*1e3, f(x), color='#845336')
plt.show()


##
## underdamped oscillations
##

R  = ufloat(200, 8)           # resistance
L  = ufloat(3.82e-3, 0.5e-03) # inductance
C  = ufloat(44.5e-9, 0.5e-9)  # capacitance
Vo = 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.318, -0.571, -0.761, -0.939, -1.088, -1.210, -1.290, -1.348, -1.379,
          -1.350, -1.299, -1.253, -1.130, -0.993, -0.873, -0.721, -0.558, -0.391,
          -0.238, -0.081,  0.075,  0.221,  0.359,  0.473,  0.573,  0.650,  0.726,
           0.760,  0.789,  0.790,  0.778,  0.740,  0.684,  0.646,  0.561,  0.484,
           0.390,  0.301,  0.201,  0.117,  0.023, -0.060, -0.149, -0.223, -0.304,
          -0.349, -0.387, -0.431, -0.437, -0.464, -0.421, -0.355, -0.260, -0.160,
          -0.052,  0.048,  0.148,  0.206,  0.249,  0.257,  0.266,  0.231,  0.168,
           0.109,  0.060, -0.009, -0.068, -0.098, -0.141, -0.147, -0.154, -0.138,
          -0.110, -0.089, -0.043,  0.002,  0.029,  0.056,  0.073,  0.079,  0.094,
           0.084,  0.064,  0.055,  0.026,  0.001, -0.014)

# 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 (discharge)')
plt.xlabel('time (μs)')
plt.ylabel('tension (V)')
plt.scatter(t*1e6, V, s=12, color='#a17e3e')
plt.plot(x*1e6, f(x), color='#daba8b')
plt.show()


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

# time
t = array(0,  10,  20,  30,  40,  60,  80, 100, 120,  160,  220,  280,
        340, 420, 500, 560, 620, 700, 780, 860, 940, 1040, 1120, 1180,
       1330)*1e-7

# tension
V = -array(0.18, 0.91, 1.54, 2.03, 2.52, 3.11, 3.58, 3.70, 3.77, 3.59, 3.04, 2.37,
           1.76, 1.10, 0.72, 0.49, 0.30, 0.19, 0.07, 0.08, 0.05, 0.02, 0.01, 0.03,
           0.01)


## 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.01, guess=[a.n, y.n])

# χ² and compatibility tests
alpha = chi_squared_fit(t, V, f, sigma=0.01, 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 (discharge)')
plt.xlabel('time (μs)')
plt.ylabel('tension (V)')
plt.scatter(t*1e6, V, color='#5c6652')
plt.plot(x*1e6, f(x), color='#718062')


## 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='#8c4f4a')
plt.plot(x*1e3, f(x), color='#845336')

## 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='#daba8b')
plt.show()