# 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 capacitor (II)
## (all SI units)

C = ufloat(8.43e-7, 1e-8) # capacitor
R = ufloat(996, 4)        # resistor

# frequency
nu = array(60, 150, 200, 300, 400, 500, 700, 800,
           1e3, 1.5e3, 2e3, 2.5e3, 3e3, 4e3, 5e3, 10e3) 

# V input
Va = array(10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0,
           10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 9.92)/2

# V output
Vb = array(2.40, 4.32, 5.02, 6.12, 6.86, 7.30, 7.90, 8.28,
           8.55, 9.24, 9.30, 9.40, 9.60, 9.62, 9.65, 9.70)/2

# time offset Va - Vb
Oab = array(3.4e-3, 980e-6, 602e-6, 295e-6, 176e-6, 116e-6,  63e-6,  46e-6,
             31e-6,  12e-6,   8e-6, 4.8e-6, 4.1e-6, 2.1e-6, 1.2e-6, 0.4e-6)*(-1)

# time offset I - V
Oiv = array(4.2e-3, 1.71e-3, 1.26e-3, 820e-6, 630e-6, 500e-6, 348e-6, 318e-6,
            248e-6,  164e-6,  125e-6, 101e-6,  83e-6,  62e-6,  50e-6,  25e-6)*(-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*Co.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, g(x))
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))