# 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 *

##
## Impedance of a resistor
## (all SI units)

R = ufloat(996, 4) # resistor (for oscilloscope)

# frequency
nu = array(60, 100, 150, 200, 300, 500, 750, 1e3, 1.5e3,  2e3,  3e3,
           5e3, 8e3, 12e3, 15e3, 20e3, 25e3, 40e3, 50e3, 60e3, 75e3,
           100e3, 150e3, 200e3, 250e3, 300e3, 350e3)

# tension drop (via voltmeter)
Vm = array(1.564, 1.564, 1.562, 1.558, 1.563, 1.552, 1.536, 1.518,
           1.470, 1.417, 1.303, 1.096, 0.878, 0.715, 0.644, 0.575,
           0.533, 0.470, 0.445, 0.423, 0.395, 0.349, 0.257, 0.169,
           0.085, 0.038, 0.020)*np.sqrt(2)

# current (via amperometer)
Im = array(1.663, 1.664, 1.663, 1.664, 1.664, 1.664, 1.664, 1.664,
           1.664, 1.664, 1.664, 1.666, 1.671, 1.679, 1.687, 1.702,
           1.721, 1.798, 1.865, 1.943, 2.076, 2.327, 2.817, 3.143,
           3.255, 3.294, 3.301)*1e-3*np.sqrt(2)

# tension input - ground (via oscilloscope)
Va = array(9.6, 9.6, 9.6, 9.6, 9.6, 9.6, 9.6, 9.6, 9.6, 9.6,
           9.6, 9.6, 9.6, 9.6, 9.6, 9.6, 9.6, 9.6, 9.6, 9.6,
           9.6, 9.6, 9.6, 9.6, 9.6, 9.6, 9.6)/2

# tension output - ground (via oscilloscope)
Vb = array(4.8, 4.8, 4.8, 4.8, 4.8, 4.8, 4.8, 4.8, 4.8, 4.8,
           4.8, 4.8, 4.8, 4.8, 4.8, 4.8, 4.8, 4.8, 4.8, 4.8,
           4.8, 4.8, 4.8, 4.8, 4.8, 4.8, 4.8)/2

# add uncertainties
Vb, Va = ufloat(Vb[0], 0.01), ufloat(Va[0], 0.01)

Io = Vb/R     # current (via oscilloscope)
Vo = Va - Vb  # tension drop (via oscilloscope)
Z  = Vo/Io    # impedance

plt.figure(1)
plt.clf()
plt.xlabel("frequency (Hz)")
plt.ylabel("voltage voltmeter / oscilloscope")
plt.semilogx(nu, Vm/Vo.n, 'o-', color='#bb3e5f', markersize=4.5)
plt.show()

plt.figure(2)
plt.clf()
plt.xlabel("frequency (Hz)")
plt.ylabel("current amperometer / oscilloscope")
plt.semilogx(nu, Io.n/Im, 'o-', color="#56ad5e", markersize=4.5)
plt.show()

plt.figure(3)
plt.clf()
plt.title('impedance (purely resistive circuit)')
plt.xlabel('frequency (Hz)')
plt.ylabel('magnitude (kΩ)')
plt.semilogx(nu, nu/nu*Z.n/1e3, 'o-', color="#d0aa23", markersize=4.5)
plt.show()

alpha = check_measures(Z, R)

print(mformat('''
Z: {:.3}Ω
R: {:.3}Ω

compatibility test:
  α={:.2f}, α>ε: {}
''', Z, R, alpha, alpha>epsilon))


##
## Impedance of a capacitor
## (all SI units)

C = ufloat(10.74e-9, 1e-9) # capacitor
R = ufloat(996, 4)         # resistor

# frequency (Hz)
nu = array(60, 75, 100, 150, 200, 300, 400, 500, 600, 700, 800, 1000, 1200,
           1300, 1400, 1500, 1800, 2500, 5000, 10e3, 50e3, 100e3, 150e3,
           200e3, 300e3, 350e3)

# V input
Va = array(10.2, 10.2, 10.2, 10.2, 10.2, 10.2, 10.2, 10.2, 10.2, 10.2, 10.2,
           10.2, 10.2, 10.2, 10.2, 10.2, 10.2, 10.2, 10.2, 10.2, 10.2, 10.2,
           10.2, 10.2, 10.2, 10.2)/2

# V output
Vb = array(45.0e-3, 55.2e-3,  74e-3, 106e-3, 140e-3, 212e-3, 276e-3,
            344e-3,  460e-3, 520e-3, 580e-3, 700e-3, 860e-3, 920e-3,
            1.05,  1.16,  1.41, 1.88, 3.60, 6.15, 9.89, 10.09, 10.16,
           10.16, 10.17, 10.17)/2

# time offset Va - Vb
Oab = array(4.123,   3.300,    2.480,   1.650,   1.230,   0.820,   0.611,   0.490,
            0.404,   0.344,    0.302,   0.240,   0.196,   0.181,   0.167,   0.155,
            0.128,   0.089,    0.040,   0.016,   0.95e-3, 0.23e-3, 0.10e-3,
            0.05e-3, 0.025e-3, 0.020e-3)*(-1e-3)

# time offset I - V
Oiv = array(4.200e-3, 3.320e-3, 2.500e-3, 1.645e-3, 1.250e-3,  830e-6,  635e-6,
              498e-6,   420e-6,   350e-6,   310e-6,   250e-6,  208e-6,  192e-6,
              176e-6,   166e-6,   138e-6,   100e-6,    50e-6,   25e-6, 4.91e-6,
             2.49e-6,  1.68e-6,  1230e-9,   826e-9,   720e-9)*(-1)

# estimate uncertainties
Evb, Eva = ufloat(Vb[0], 1e-3), ufloat(Va[0], 0.1)
sigmaZ   = (R*(Eva-Evb)/Evb).s
sigmaH1  = ((Eva-Evb)/Evb).s
sigmaH2  = ((Evb-Eva)/Evb).s

om  = 2*np.pi*nu # angular frequency
I   = Vb/R.n     # output current 
Vab = Va - Vb    # tension drop
Fab = om * Oab   # phase difference Vᵢ - Vₒ
Fiv = om * Oiv   # phase difference I - Vₒ

Z  = Vab/I  * np.exp(1j*Fiv) # impedance
H1 = Vab/Vb * np.exp(1j*Fab) # transfer function ΔV→Vb
H2 = Vb/Va  * np.exp(1j*Fab) # transfer function Va→Vb

# plot and fit Z(ν)
plt.figure(4)
plt.clf()

# magnitude
plt.subplot(2, 1, 1)
plt.title('impedance (capacitive circuit)')
plt.ylabel('magnitude (kΩ)')
plt.semilogx(nu, abs(Z)/1e3, 'o', color="#36913d", markersize=4.5)

# fit Y=kX where Y=|Z|, X=1/ν, k=1/2πC
k  = simple_linear(1/nu, abs(Z), sigmaZ)
Co = 1/(2*np.pi*k)
f  = lambda x: k.n/x

x = np.arange(nu.min()-10, nu.max(), 10)
plt.semilogx(x, f(x)/1e3, color='#589f22')

# phase
plt.subplot(2, 1, 2)
plt.xlabel('frequency (Hz)')
plt.ylabel('phase (rad)')
plt.ylim(-1.8, -1)
plt.semilogx(nu, Fiv, 'o', color="#36913d", markersize=4.5)
plt.semilogx(x, x/x * -np.pi/2, color="#589f22")
plt.show()

phi   = sample(Fiv).val()
alpha = check_measures(phi, ufloat(-np.pi/2, 0))

print(mformat('''
φ:    {} rad
-π/2: {:.4} rad

compatibility test:
  α={:.2f}, α>ε: {}
''', phi, -np.pi/2, alpha, alpha>epsilon))

alpha = check_measures(C, Co)
beta  = chi_squared_fit(nu, abs(Z), f, sigmaZ)

print(mformat('''
k: {}
C: {} F
Cₒ: {} F

compatibility test:
  α={:.2f}, α>ε: {}

χ² test:
  β={:.2f}, β>ε: {}
''', k, C, Co,
     alpha, alpha>epsilon,
     beta, beta>epsilon))


# plot, fit H₁(ν)
plt.figure(5)
plt.clf()

# magnitude
plt.subplot(2,1,1)
plt.title('transmission function 1')
plt.ylabel('amplitude (Vout-Vin / Vout)')
plt.semilogx(nu, abs(H1), 'o', color="#9b2e83", markersize=4.5)

# fit Y=kX where Y=|H1|, X=1/ν, k=1/2πRC
k   = simple_linear(1/nu, abs(H1), sigmaH1)
RCo = 1/(2*np.pi*k)
f   = lambda x: k.n/x

x = np.arange(nu.min()-10, nu.max(), 10)
plt.semilogx(x, f(x), color='#9b2e83')

# phase
plt.subplot(2,1,2)
plt.xlabel('frequency (Hz)')
plt.ylabel('phase (rad)')
plt.semilogx(nu, Fab, 'o', color="#3a44ad", markersize=4.5)
plt.semilogx(x, -np.pi/2+np.arctan(2*np.pi*x*R.n*C.n))
plt.show()

alpha = check_measures(R*C, RCo)
beta  = chi_squared_fit(nu, abs(H1), f, sigmaH1)

print(mformat('''
k:   {} Hz
RC:  {} s
RCₒ: {} s

compatibility test:
  α={:.2f}, α>ε: {}

χ² test:
  β={:.2f}, β>ε: {}
''', k, R*C, RCo,
     alpha, alpha>epsilon,
     beta, beta>epsilon))


# plot, fit H₂(ν)
plt.figure(6)
plt.clf()

# magnitude
plt.subplot(2,1,1)
plt.title('transmission function 2')
plt.ylabel('amplitude (Vout / Vin)')
plt.semilogx(nu, abs(H2), 'o', color="#9b2e83", markersize=4.5)

# fit Y=a+bX where Y=1/|H₂|², X=1/ν², a=1, b=1/(2πRC)²
a,b = linear(1/nu**2, 1/abs(H2)**2, 0.01)
RCo = 1/(2*np.pi*um.sqrt(b))
f   = lambda x: 1/np.sqrt(1 + b.n/x**2)

x = np.arange(nu.min()-10, nu.max(), 10)
plt.semilogx(x, f(x), color='#9b2e83')

# phase
plt.subplot(2,1,2)
plt.xlabel('frequency (Hz)')
plt.ylabel('phase (rad)')
plt.semilogx(nu, Fab, 'o', color="#3a44ad", markersize=4.5)
plt.semilogx(x, -np.pi/2+np.arctan(2*np.pi*x*R.n*C.n))
plt.show()

alpha = check_measures(R*C, RCo)
beta  = chi_squared_fit(nu, abs(H2), f, sigmaH2)

print(mformat('''
b:   {}
RC:  {} s
RCₒ: {} s

compatibility test:
  α={:.2f}, α>ε: {}

χ² test:
  β={:.2f}, β>ε: {}
''', b, R*C, RCo,
     alpha, alpha>epsilon,
     beta, beta>epsilon))