import numpy as np
from qiskit_dynamics.models import (
    HamiltonianModel,
    LindbladModel,
    rotating_wave_approximation,
)
from qiskit_dynamics.signals import Signal

Omega_0 = 0.5
tau = 3
omega_0 = 1
nu = 1/(2*np.pi) #2\pi\nu = \omega_0

Omega = lambda t: Omega_0 * np.exp(-t**2/(2*tau**2))
drive_signal = Signal(Omega,nu,0)

drive_operator = np.array([[[0,-1],[-1,0]]])
drift_operator = omega_0*np.array([[1,0],[0,-1]])

model = HamiltonianModel(drive_operator,drift=drift_operator,signals=[drive_signal])

print(model.evaluate(0))

model.set_evaluation_mode("sparse")

print(model.evaluate(0))

model.rotating_frame = np.array([1,-1])
print(model(0))

rwa_model = rotating_wave_approximation(model, cutoff_freq=1/4)
print(np.round(rwa_model(0),6))

model.rotating_frame = np.array([1,-1])

dis_op = np.array([[[0,0],[1,0]]])
rho = np.diag([1,0]) # excited state
dis_model = LindbladModel.from_hamiltonian(model, dissipator_operators=dis_op, dissipator_signals=[1], evaluation_mode="dense")

print(help(dis_model.set_evaluation_mode))

dis_model.set_evaluation_mode('dense')
print(dis_model.evaluate_hamiltonian(0))
print(dis_model(0,rho))

dis_model.set_evaluation_mode('sparse')
print(dis_model.evaluate_hamiltonian(0))
print(dis_model(0,rho))

rho_vec = rho.flatten(order="F") # vectorized operator using column stacking convention
dis_model.set_evaluation_mode('dense_vectorized')
print(dis_model(0))
print(dis_model(0) @ rho_vec)
print(dis_model(0,rho_vec))

sig_vals = model.signals(0)

model.set_evaluation_mode("dense")
operator_collection = model._operator_collection
print(type(operator_collection))
print(operator_collection.evaluate(sig_vals))

model.set_evaluation_mode("sparse")
operator_collection = model._operator_collection
print(type(operator_collection))
print(operator_collection.evaluate(sig_vals))