import sys, cmath
import numpy as np
import matplotlib.pyplot as plt

from pscan2 import psconfig

tjm_coeff = psconfig.tjm_models['tjm1']

A = np.array(tjm_coeff[0])
B = np.array(tjm_coeff[1])
P = np.array(tjm_coeff[2])

N = len(A)

print("N: ",N, "\nA: ", A, "\nB: ", B, "\nP: ", P)

W = np.arange(0.0, 3.0, 0.01)
Nw = len(W)

JP = np.array(np.zeros(Nw), dtype=complex)
JQ = np.array(np.zeros(Nw), dtype=complex)

for i in range(0,Nw):
    w = W[i]
    iw = 1j * w
    for j in range(0,N):
        JP[i] += A[j] /(iw - P[j]) + A[j].conjugate() / (iw - P[j].conjugate())
        JQ[i] += B[j] /(iw - P[j]) + B[j].conjugate() / (iw - P[j].conjugate())
    JP[i] *= 0.5
    JQ[i] = complex(JQ[i].real/2.0,w - JQ[i].imag/2.0)

plt.plot(W, JP.real, linestyle='-', marker='*', color = 'blue')
plt.plot(W, JP.imag, linestyle='-', marker='*', color = 'red')
plt.plot(W, JQ.real, linestyle='-', marker='*', color = 'green')
plt.plot(W, JQ.imag, linestyle='-', marker='*', color = 'orange')

plt.grid(True)
plt.show()


