# coding: utf-8
from __future__ import print_function

from sympy import *
from sympy.utilities.lambdify import lambdify
import matplotlib.pyplot as plt
import numpy as np

##
## Transfer function calculator
##

I # imaginary unit
i  = Symbol('i',  positive=True, real=True) # current
R  = Symbol('R',  positive=True, real=True) # resistance
Rl = Symbol('Rl', positive=True, real=True) # resistance
L  = Symbol('L',  positive=True, real=True) # inductance
C  = Symbol('C',  positive=True, real=True) # capacitance
w  = Symbol('w',  positive=True, real=True) # angular frequency

# impedances
Zl = Rl+I*w*L
Zc = 1/(I*w*C)
Zr = R

# input/output tensions
Va = (Zc+Zr+Zl)*i
Vb = (Zc)*i

# transfer function
H = Vb/Va
phase = simplify(atan2(re(H), im(H)))
magn  = simplify(sqrt(re(H)**2 + im(H)**2))

print('arg H =', phase)
print('  |H| =', magn)


##
## Plot function
##

p = (9.94e3, 20, 3.8e-3, 4.45e-8) # R,Rl,L,C
x = np.linspace(100, 150e3, 1500)
f = np.vectorize(lambda nu: lambdify((w, R, Rl, L, C), magn)(2*pi*nu, *p))
g = np.vectorize(lambda nu: lambdify((w, R, Rl, L, C), phase)(2*pi*nu, *p))

plt.subplot(2, 1, 1)
plt.title('transfer function')
plt.xscale('log')
plt.ylabel('magnitude')
plt.plot(x, f(x))

plt.subplot(2, 1, 2)
plt.xscale('log')
plt.plot(x, g(x))
plt.xlabel('frequency (Hz)')
plt.ylabel('phase (rad)')
plt.show()