from matplotlib import pyplot as plt
import numpy as np
from scipy.integrate import quad
from qiskit.quantum_info.states import Statevector, DensityMatrix
from qiskit.quantum_info.operators import Operator
from qiskit.quantum_info.operators import SuperOp, Choi, Kraus, PTM

from qiskit_dynamics import solve_ode, solve_lmde
from qiskit_dynamics.signals import Signal
from qiskit_dynamics.models import (HamiltonianModel, 
                                    LindbladModel, 
                                    rotating_wave_approximation)


def gaussian(amp, sig, t0, t):
    return amp * np.exp( -(t - t0)**2 / (2 * sig**2) )

amp = 1. # amplitude
sig = 2. # sigma
t0 = 3.5*sig # center of Gaussian
T = 7*sig # end of signal

gaussian_envelope = lambda t: gaussian(amp, sig, t0, t)

gauss_signal = Signal(envelope=gaussian_envelope, carrier_freq=0.5)

print(gauss_signal.envelope(0.25))
print(gauss_signal(0.25))

gauss_signal.draw(0, T, 100, function='envelope')

gauss_signal.draw(0, T, 200)

#####################
# construct operators
#####################
r = 0.5
w = 1.
X = Operator.from_label('X')
Y = Operator.from_label('Y')
Z = Operator.from_label('Z')

drift = 2 * np.pi * w * Z/2
operators = [2 * np.pi * r * X/2]

###################
# construct signals
###################

# Define gaussian envelope function to have max amp and area approx 2 
amp = 1.
sig = 0.399128/r
t0 = 3.5*sig
T = 7*sig
gaussian_envelope = lambda t: gaussian(amp, sig, t0, t)

signals = [Signal(envelope=gaussian_envelope, carrier_freq=w)]

#################
# construct model
#################

hamiltonian = HamiltonianModel(drift=drift, operators=operators, signals=signals)

print(hamiltonian.evaluate(0.12))

hamiltonian.get_drift()

def plot_qubit_hamiltonian_components(hamiltonian, t0, tf, N=200):
    t_vals = np.linspace(t0, tf, N)
    model_vals = np.array([1j*hamiltonian.evaluate(t) for t in t_vals])
    x_coeff = model_vals[:, 0, 1].real
    y_coeff = -model_vals[:, 0, 1].imag
    z_coeff = model_vals[:, 0, 0].real
    plt.plot(t_vals, x_coeff, label='X component')
    plt.plot(t_vals, y_coeff, label='Y component')
    plt.plot(t_vals, z_coeff, label='Z component')
    plt.legend()

plot_qubit_hamiltonian_components(hamiltonian, 0., T)

hamiltonian.rotating_frame = hamiltonian.get_drift()

print(hamiltonian.evaluate(0.12))

# validate with independent computation
t = 0.12
-1j*(2 * np.pi * r * np.real(signals[0](t)) * (np.cos(2*np.pi * w * t) * X / 2 
                                                - np.sin(2*np.pi * w * t) * Y / 2 ))

plot_qubit_hamiltonian_components(hamiltonian, 0., T)

# set the cutoff frequency
rwa_hamiltonian = rotating_wave_approximation(hamiltonian, 2*w)

# evaluate again
print(rwa_hamiltonian.evaluate(0.12))

plot_qubit_hamiltonian_components(rwa_hamiltonian, 0., T)

# solve the problem, with some options specified
y0 = Statevector([0., 1.])
%time sol = solve_lmde(hamiltonian, t_span=[0., T], y0=y0, atol=1e-10, rtol=1e-10)
y = sol.y[-1]
print("\nFinal state:")
print("----------------------------")
print(y)
print("\nPopulation in excited state:")
print("----------------------------")
print(np.abs(y[0])**2)

%time sol = solve_lmde(rwa_hamiltonian, t_span=[0., T], y0=y0, atol=1e-10, rtol=1e-10)
y = sol.y[-1]
print("\nFinal state:")
print("----------------------------")
print(y)
print("\nPopulation in excited state:")
print("----------------------------")
print(np.abs(y[0])**2)

# construct quantum model with noise operators
noise_ops = [np.array([[0., 0.], 
                       [1., 0.]])]
noise_signals = [0.001]

lindblad_model = LindbladModel.from_hamiltonian(hamiltonian=hamiltonian, 
                                                dissipator_operators=noise_ops,
                                                dissipator_signals=noise_signals)

lindblad_model = LindbladModel(hamiltonian_operators=operators,
                               drift=drift,
                               rotating_frame=drift,
                               hamiltonian_signals=signals,
                               dissipator_operators=noise_ops,
                               dissipator_signals=noise_signals)



# density matrix
y0 = DensityMatrix([[0., 0.], [0., 1.]])

%time sol = solve_lmde(lindblad_model, t_span=[0., T], y0=y0, atol=1e-10, rtol=1e-10)
sol.y[-1]

rwa_lindblad_model = rotating_wave_approximation(lindblad_model, 2.0*w)
%time sol = solve_lmde(rwa_lindblad_model, t_span=[0., T], y0=y0, atol=1e-10, rtol=1e-10)
sol.y[-1]

lindblad_model.evaluation_mode = 'dense_vectorized'
lindblad_model.rotating_frame = None

# identity quantum channel in superop representation
y0 = SuperOp(np.eye(4))

%time sol = solve_lmde(lindblad_model.__call__, t_span=[0., T], y0=y0, atol=1e-12, rtol=1e-12)
print(sol.y[-1])

print(PTM(sol.y[-1]))

