%matplotlib inline
import matplotlib.pyplot as plt

import numpy as np

from sympy import *
init_printing()

from sympsi import *
from sympsi.boson import *
from sympsi.pauli import *
from sympsi.operatorordering import *
from sympsi.expectation import *
from sympsi.operator import OperatorFunction

w, t, Nth, Ad, kappa = symbols(r"\omega, t, n_{th}, A_d, kappa", positive=True)

a = BosonOp("a")
rho = Operator(r"\rho")
rho_t = OperatorFunction(rho, t)

H = w * Dagger(a) * a + Ad * (a + Dagger(a))

Eq(Symbol("H"), H)

c_ops = [sqrt(kappa * (Nth + 1)) * a, sqrt(kappa * Nth) * Dagger(a)]

me = master_equation(rho_t, t, H, c_ops)

me

# first setup time-dependent operators
a_t = OperatorFunction(a, t)
a_to_a_t = {a: a_t, Dagger(a): Dagger(a_t)}
H_t = H.subs(a_to_a_t)
c_ops_t = [c.subs(a_to_a_t) for c in c_ops]

# operator master equation for a
ome_a = operator_master_equation(a_t, t, H_t, c_ops_t)

Eq(ome_a.lhs, normal_ordered_form(ome_a.rhs.doit().expand()))

# operator master equation for n = Dagger(a) * a
ome_n = operator_master_equation(Dagger(a_t) * a_t, t, H_t, c_ops_t)

Eq(ome_n.lhs, normal_ordered_form(ome_n.rhs.doit().expand()))

ops, op_eqm, sc_eqm, sc_ode, ofm, oim = semi_classical_eqm(H, c_ops)

html_table([[Eq(Expectation(key), ofm[key]), sc_ode[key]] for key in operator_sort_by_order(sc_ode)])

A_eq, A, M, b = semi_classical_eqm_matrix_form(sc_ode, t, ofm)

A_eq

A_sol = M.LUsolve(-b)

A_sol[ops.index(Dagger(a)*a)]

A_sol[ops.index(a)]

A_sol[ops.index(Dagger(a))]

solve([eq.rhs for eq in sc_ode.values()], list(ofm.values()))

sols = dsolve(list(sc_ode.values())); sols

# hack
tt = [s for s in sols[0].rhs.free_symbols if s.name == 't'][0]

ics = {ofm[Dagger(a)].subs(tt, 0): 2, 
       ofm[a].subs(tt, 0): 2,
       ofm[Dagger(a)*a].subs(tt, 0): 4}; ics

constants = set(sum([[s for s in sol.free_symbols if (str(s)[0] == 'C')] for sol in sols], [])); constants

C_sols = solve([sol.subs(tt, 0).subs(ics) for sol in sols], constants); C_sols

sols_with_ics = [sol.subs(C_sols) for sol in sols]; sols_with_ics

values = {w: 1.0, Ad: 0.0, kappa: 0.1, Nth: 0.0}

sols_funcs = [sol.rhs.subs(values) for sol in sols_with_ics]; sols_funcs

times = np.linspace(0, 50, 500)

y_funcs = [lambdify([tt], sol_func, 'numpy') for sol_func in sols_funcs]

fig, axes = plt.subplots(len(y_funcs), 1, figsize=(12, 6))

for n, y_func in enumerate(y_funcs):
    axes[n].plot(times, np.real(y_func(times)), 'r')
    
axes[2].set_ylim(0, 5);

sx, sy, sz, sm, sp = SigmaX(), SigmaY(), SigmaZ(), SigmaMinus(), SigmaPlus()

Omega, gamma_0, N, t = symbols("\Omega, \gamma_0, N, t", positive=True)

values = {Omega: 1.0, gamma_0: 0.5, N: 1.75}

H = -Omega/2 * sx
H

c_ops = [sqrt(gamma_0 * (N + 1)) * pauli_represent_x_y(sm), 
         sqrt(gamma_0 * N) * pauli_represent_x_y(sp)]

ops, op_eqm, sc_eqm, sc_ode, ofm, oim = semi_classical_eqm(H, c_ops)

html_table([[Eq(Expectation(key), ofm[key]), sc_ode[key]] for key in operator_sort_by_order(sc_ode)])

A_eq, A, M, b = semi_classical_eqm_matrix_form(sc_ode, t, ofm)

A_eq

A_sol = M.LUsolve(-b)

A_sol[ops.index(sx)]

A_sol[ops.index(sy)]

A_sol[ops.index(sz)]

pauli_represent_x_y(sp).subs({sx: A_sol[ops.index(sx)], sy: A_sol[ops.index(sy)]})

pauli_represent_x_y(sm).subs({sx: A_sol[ops.index(sx)], sy: A_sol[ops.index(sy)]})

solve([eq.rhs for eq in sc_ode.values()], list(ofm.values()))

solve([eq.subs(N, 0).rhs for eq in sc_ode.values()], list(ofm.values()))

%reload_ext version_information

%version_information sympy, sympsi