import numpy as np
import matplotlib.pyplot as plt

wr  = 7.10      # GHz
wq  = 5.20      # GHz
g   = 0.120     # GHz
a   = -0.250    # GHz, transmon anharmonicity
kap = 0.0012    # GHz
n   = 8

d     = wq - wr
chi   = g*g*a/(d*(d + a))
ncrit = d*d/(4*g*g)
fg    = wr - chi
fe    = wr + chi
xs    = np.linspace(wr - 0.015, wr + 0.015, 600)

def lor(f):
    return 1/np.sqrt(1 + (2*(xs - f)/kap)**2)

def ptr(f):
    return np.sqrt(n)*(kap/2)/(kap/2 + 1j*(wr - f))

ag, ae = ptr(fg), ptr(fe)

print(f"Delta/2pi = {d*1e3:.1f} MHz")
print(f"chi/2pi   = {chi*1e3:.2f} MHz")
print(f"2chi/kappa= {2*abs(chi)/kap:.2f}")
print(f"ncrit     = {ncrit:.1f}")

fig, ax = plt.subplots(1, 2, figsize=(8, 3.2), layout="constrained")
ax[0].plot(xs, lor(fg), label="qubit |g>")
ax[0].plot(xs, lor(fe), label="qubit |e>")
ax[0].set(xlabel="probe frequency (GHz)", ylabel="|S21|")
ax[0].legend()
ax[1].scatter([ag.real, ae.real], [ag.imag, ae.imag])
ax[1].plot([ag.real, ae.real], [ag.imag, ae.imag])
ax[1].set(xlabel="I", ylabel="Q", aspect="equal")
plt.show()