# To use the package locally, add the C2QA repository's root folder to the path prior to importing c2qa.
import os
import sys
module_path = os.path.abspath(os.path.join("../.."))
if module_path not in sys.path:
    sys.path.append(module_path)
    
import c2qa
import numpy as np
import qiskit

# Define Hamiltonian parameters
omega_R = 2
omega_Q = 5
chi = 0.1

from qiskit.circuit import Parameter
# Set number of qubits per qumode
num_qubits_per_qumode = 3

# Define parameter dt
dt = Parameter('dt')

# Create new circuit
qmr = c2qa.QumodeRegister(num_qumodes=1, num_qubits_per_qumode=num_qubits_per_qumode)
qbr = qiskit.QuantumRegister(1)
U_JC = c2qa.CVCircuit(qmr,qbr)

# Append U_R
U_JC.cv_r(-omega_R*dt,qmr[0])

# Append U_Q
U_JC.rz(omega_Q*dt,qbr[0])

# Append U_\chi -- KS: this needs to be updated to reflect naming conventions in manuscript
U_JC.cv_c_r(-chi*dt/2,qmr[0],qbr[0])

# Compile this circuit into a single parameterized gate
U_JC = U_JC.to_gate(label='U_JC')

qmr = c2qa.QumodeRegister(num_qumodes=1, num_qubits_per_qumode=num_qubits_per_qumode)
qbr = qiskit.QuantumRegister(1)

# Instantiate the circuit and initialize the qubit to the '+' state.
circuit = c2qa.CVCircuit(qmr,qbr)
circuit.initialize([1/np.sqrt(2), 1/np.sqrt(2)], qbr)

# Append the time evolution unitary
circuit.append(U_JC,qmr[0] + [qbr[0]]);

# Draw the circuit
circuit.decompose().draw(output='mpl')

I_qumode = circuit.ops.get_eye(qmr.cutoff).toarray()

sigmax = np.kron(np.array([[0, 1], [1, 0]]), I_qumode)
sigmaz = np.kron(np.array([[1, 0], [0, -1]]), I_qumode)

# First construct operator sigma^x. First, we need to retrieve the identity operator over the qumode. As its dimensionality is
# dependent upon the cutoff, the easiest way to do this is to use

tarray = np.linspace(0,3*2*np.pi/omega_Q,100)
sigma_x_expectation = []
sigma_z_expectation = []
for tau in tarray:
    state, result, _ = c2qa.util.simulate(circuit.assign_parameters({dt : tau}))
    sigma_x_expectation.append(state.expectation_value(sigmax).real)
    sigma_z_expectation.append(state.expectation_value(sigmaz).real)

import matplotlib.pyplot as plt

fontsize = 20

plt.figure(figsize=(15,4))
plt.subplot(1,2,1)
plt.plot(tarray/(2*np.pi/omega_Q),sigma_x_expectation,linewidth=2)
plt.ylabel(r'$\langle\sigma^x\rangle$',fontsize=fontsize,usetex=True);
plt.xlabel(r'$t$ $(2\pi/\omega_Q)$',fontsize=fontsize,usetex=True);
plt.xlim([0,3]);
plt.ylim([-1.05,1.05]);

plt.subplot(1,2,2)
plt.plot(tarray/(2*np.pi/omega_Q),sigma_z_expectation,linewidth=2,color='tab:red')
plt.ylabel(r'$\langle\sigma^z\rangle$',fontsize=fontsize,usetex=True);
plt.xlabel(r'$t$ $(2\pi/\omega_Q)$',fontsize=fontsize,usetex=True);
plt.xlim([0,3]);
plt.ylim([-1.05,1.05]);

fockNumberList = [0,1,2,3,4]
tarray = np.linspace(0,3*2*np.pi/omega_Q,100)

plt.figure(figsize=(15,4))
plt.ylabel(r'$\langle\sigma^x\rangle$',fontsize=fontsize,usetex=True);
plt.xlabel(r'$t$ $(2\pi/\omega_Q)$',fontsize=fontsize,usetex=True);
plt.xlim([0,3]);
plt.ylim([-1.05,1.05]);

for fockNumber in fockNumberList:
    print("Time evolving for Fock state |" + str(fockNumber) + '>...')
    qmr = c2qa.QumodeRegister(num_qumodes=1, num_qubits_per_qumode=num_qubits_per_qumode)
    qbr = qiskit.QuantumRegister(1)

    # Instantiate the circuit and initialize the qubit to the '+' state.
    circuit = c2qa.CVCircuit(qmr,qbr)
    circuit.initialize([1/np.sqrt(2), 1/np.sqrt(2)], qbr)

    # Initialize the qumode
    circuit.cv_initialize(fockNumber,qmr[0])

    circuit.append(U_JC,qmr[0] + [qbr[0]]);
    
    sigma_x_expectation = []
    for tau in tarray:
        state, result, _ = c2qa.util.simulate(circuit.assign_parameters({dt : tau}))
        sigma_x_expectation.append(state.expectation_value(sigmax).real)


    plt.plot(tarray/(2*np.pi/omega_Q),sigma_x_expectation,linewidth=2, label = r'$|' + str(fockNumber) + r'\rangle $')

plt.legend()

# Choose alpha for coherent state
alpha = 1

# Choose total animation time
total_time = 1*2*np.pi/omega_R

# First construct circuit_qubit0
qmr = c2qa.QumodeRegister(num_qumodes=1, num_qubits_per_qumode=num_qubits_per_qumode)
qbr = qiskit.QuantumRegister(1)

# Create new circuit
qmr = c2qa.QumodeRegister(num_qumodes=1, num_qubits_per_qumode=num_qubits_per_qumode)
qbr = qiskit.QuantumRegister(1)
U_JC = c2qa.CVCircuit(qmr,qbr)
# Append U_R
U_JC.cv_r(-omega_R*total_time,qmr[0])
# Append U_Q
U_JC.rz(omega_Q*total_time,qbr[0])
# Append U_\chi -- KS: this needs to be updated to reflect naming conventions in manuscript
U_JC.cv_c_r(-chi*total_time/2,qmr[0],qbr[0])
# Compile this circuit into a single parameterized gate
U_JC = U_JC.to_gate(label='U_JC')

# Instantiate the circuit and initialize the qubit to the '0' state.
circuit_0 = c2qa.CVCircuit(qmr,qbr)
circuit_0.initialize([1,0], qbr)

# Now initialize the qumode in a coherent state
cutoff = 2**num_qubits_per_qumode
coeffs = [np.exp(-np.abs(alpha)**2/2)*alpha**n/(np.sqrt(np.math.factorial(n))) for n in range(0,cutoff)]
circuit_0.cv_initialize(coeffs,qmr[0])


# Append time evolution unitary
circuit_0.append(U_JC,qmr[0] + [qbr[0]]);
# circuit_0.assign_parameters({dt : total_time})


# dt = total_time
# # Append U_R
# circuit_0.cv_r(-omega_R*dt,qmr[0])
# # Append U_Q
# circuit_0.rz(omega_Q*dt,qbr[0])
# # Append U_\chi -- KS: this needs to be updated to reflect naming conventions in manuscript
# circuit_0.cv_c_r(-chi*dt/2,qmr[0],qbr[0])

# Compile this circuit into a single parameterized gate
# U_JC = U_JC.to_gate(label='U_JC')



# # Now repeat the above steps for a qubit initialized to the '1' state:
# circuit_1 = c2qa.CVCircuit(qmr,qbr)
# circuit_1.initialize([0,1], qbr)
# circuit_1.cv_d(alpha,qmr[0])
# circuit_1.append(U_JC,qmr[0] + [qbr[0]]);
# circuit_1 = circuit_1.assign_parameters({dt : total_time})

# Animate wigner function of each circuit
c2qa.animate.animate_wigner(circuit_0, animation_segments = 1000)

