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

##
## Impedence of a capacitor (I)
## (all SI units)

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

# frequency
nu = array(60, 150, 300, 500, 800, 1e3, 2e3, 3e3, 5e3,
           10e3, 25e3, 50e3, 100e3, 125e3, 175e3, 250e3, 300e3)

# V input
Va = array(10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0,
           10.0, 9.80, 9.60, 9.60, 9.44, 9.60, 9.44, 9.44)/2

# V output
Vb = array(42.0e-3, 102e-3, 198e-3, 336e-3, 540e-3, 680e-3, 1.36, 2.0, 3.2,
           5.48, 8.20, 9.0, 9.2, 9.28, 9.40, 9.28, 9.28)/2

# time offset Va - Vb
Oab = array(4.200e-3, 1.660e-3, 820e-6, 490e-6, 300e-6, 240e-6, 114e-6, 73e-6, 39e-6,
            15.60e-6, 3.4e-6, 900e-9, 240e-9, 130e-9, 90e-9, 40e-9, 20e-9)*(-1)

# time offset I - V
Oiv = array(4.200e-3, 1.645e-3, 830e-6, 498e-6, 310e-6, 250e-6, 124e-6, 84e-6, 50e-6,
            25.30e-6, 10.01e-6, 5.0e-6, 2.49e-6, 2.01e-6, 1.426e-6, 997e-9, 826e-9)*(-1)


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


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


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

# magnitude
plt.subplot(2, 1, 1)
plt.title('impedance (RC 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('transfer function 1')
plt.ylabel('magnitude (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('transfer function 2')
plt.ylabel('magnitude (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)                 # magnitude
g   = lambda x: -np.pi/2 + np.arctan(2*np.pi*x*R.n*C.n) # phase

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)
gamma = chi_squared_fit(nu, Fab, g, sigmaF)

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

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

χ² test (magnitude):
  β={:.2f}, β>ε: {}

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