import numpy as np
from qiskit.opflow import I, X, Y, Z, MatrixEvolution, PauliTrotterEvolution
from qiskit.circuit import Parameter
from qiskit import QuantumCircuit
from qiskit import Aer
from qiskit.compiler import transpile
import qc_grader



# Define which will contain the Paulis
pauli_list = [X,Y,Z,I]

# Define Paulis and add them to the list
###INSERT CODE BELOW THIS LINE


###DO NOT EDIT BELOW THIS LINE

for pauli in pauli_list:
    print(pauli, '\n')


from qc_grader.challenges.qgss_2022 import grade_lab2_ex1

grade_lab2_ex1(pauli_list)

# Define list of ladder operators
s_plus, s_minus = (X+(1j*Y))/2, (X-(1j*Y))/2
ladder_operator_list = [s_plus, s_minus]


# Define ladder operators and add the to the list
###INSERT CODE BELOW THIS LINE

###DO NOT EDIT BELOW THIS LINE

for ladder_operator in ladder_operator_list:
    print(ladder_operator, '\n')


from qc_grader.challenges.qgss_2022 import grade_lab2_ex2

grade_lab2_ex2(ladder_operator_list)

# Define list which will contain the matrices representing the Pauli operators
matrix_sigma_list = []

# Add matrix representation of Paulis to the list
###INSERT CODE BELOW THIS LINE
for op in pauli_list:
    matrix_sigma_list.append(op.to_matrix())
    

###DO NOT EDIT BELOW THIS LINE

for matrix_sigma in matrix_sigma_list:
    print(matrix_sigma, '\n')

from qc_grader.challenges.qgss_2022 import grade_lab2_ex3

grade_lab2_ex3(matrix_sigma_list)

# Define a list which will contain the circuit representation of the Paulis
circuit_sigma_list = []

# Add circuits to list
###INSERT CODE BELOW THIS LINE
for gate in pauli_list:
    circuit_sigma_list.append(gate.to_circuit())

###DO NOT EDIT BELOW THIS LINE

for circuit in circuit_sigma_list:
    print(circuit, '\n')



from qc_grader.challenges.qgss_2022 import grade_lab2_ex4

grade_lab2_ex4(circuit_sigma_list)

# Define a variable theta to be a parameter with name 'theta'
theta = Parameter('theta')
# Set number of qubits to 1
qubits_count = 1
# Initialize a quantum circuit with one qubit
quantum_circuit = QuantumCircuit(qubits_count)

# Add a parametrized RX rotation on the qubit
###INSERT CODE BELOW THIS LINE
quantum_circuit.rx(theta,0)

###DO NOT EDIT BELOW THIS LINE

print(quantum_circuit)


quantum_circuit.draw()

from qc_grader.challenges.qgss_2022 import grade_lab2_ex5

grade_lab2_ex5(quantum_circuit)

# Set the value of the parameter
theta_value = np.pi

# Bind the value to the parametrized circuit
###INSERT CODE BELOW THIS LINE
quantum_circuit = quantum_circuit.bind_parameters({theta: theta_value})
###DO NOT EDIT BELOW THIS LINE

print(quantum_circuit)

from qc_grader.challenges.qgss_2022 import grade_lab2_ex6

grade_lab2_ex6(quantum_circuit)

print(I.to_matrix())

# Use the formula above to define the Hamiltonian operator
###INSERT CODE BELOW THIS LINE


###DO NOT EDIT BELOW THIS LINE
H = 1/2*((I^I)+(X^X)+(Y^Y)+(Z^Z))
# Get its matrix representation
H_matrix = H.to_matrix()

print(H_matrix)

from qc_grader.challenges.qgss_2022 import grade_lab2_ex7

grade_lab2_ex7(H_matrix)

# Define a parameter t for the time in the time evolution operator
t = Parameter('t')

# Follow the instructions above to define a time-evolution operator
###INSERT CODE BELOW THIS LINE
time_evolution_operator=(H*t).exp_i()

###DO NOT EDIT BELOW THIS LINE

print(time_evolution_operator)

from qc_grader.challenges.qgss_2022 import grade_lab2_ex8

grade_lab2_ex8(time_evolution_operator)

# Set a total time for the time evolution
evolution_time = 0.5

# Instantiate a MatrixEvolution() object to convert the time evolution operator
qc =  MatrixEvolution().convert(time_evolution_operator)
#  and bind the value for the time parameter
bound_matrix_exponentiation_circuit = qc.bind_parameters({t: evolution_time})
###INSERT CODE BELOW THIS LINE


###DO NOT EDIT BELOW THIS LINE

print(bound_matrix_exponentiation_circuit)

from qc_grader.challenges.qgss_2022 import grade_lab2_ex9

grade_lab2_ex9(bound_matrix_exponentiation_circuit)

# Define a value for the duration of the time-step
time_step_value = 0.1

# Instantiate a PauliTrotterEvolution() object and convert the time-evolution operator
# to then bind the value of the time step
###INSERT CODE BELOW THIS LINE
trotter_circuit=PauliTrotterEvolution().convert(time_evolution_operator)
bound_trotter_exponentiation_circuit = trotter_circuit.bind_parameters({t: time_step_value})

###DO NOT EDIT BELOW THIS LINE

print(bound_trotter_exponentiation_circuit)

from qc_grader.challenges.qgss_2022 import grade_lab2_ex10

grade_lab2_ex10(bound_trotter_exponentiation_circuit)

# Define the number of steps needed to reach the previously set total time-evolution
steps = int(evolution_time/time_step_value)

# Compose the operator for a Trotter step several times to generate the 
# operator for the full time-evolution
###INSERT CODE BELOW THIS LINE
trot = bound_trotter_exponentiation_circuit
U = trot
for i in range(steps-1):
    trot@=U 
    
total_time_evolution_circuit = trot

###DO NOT EDIT BELOW THIS LINE

print(total_time_evolution_circuit)


from qc_grader.challenges.qgss_2022 import grade_lab2_ex11

grade_lab2_ex11(total_time_evolution_circuit)

# Set number of qubits
num_qubits = 3
# Define time parameter
t = Parameter('t')
# Set total evolution time
evolution_time_t = 2
# Set size of time-step for Trotter evolution
time_step_value_t = 0.1
# Define the number of steps
steps_t = int(evolution_time_t/time_step_value_t)
# Create circuit
tight_binding_circuit = QuantumCircuit(num_qubits)
# Add initial state preparation
tight_binding_circuit.x(0)

# Define the Hamiltonian, the time-evolution operator, the Trotter step and the total evolution
###INSERT CODE BELOW THIS LINE
XXs = (I^X^X) + (X^X^I) #make the sigmax first term
YYs = (I^Y^Y) + (Y^Y^I) #make the sigmay second term
H = (XXs + YYs) # prepare the Hamiltonian
time_evolution_operator = (t * H).exp_i() # make the time evolution 
trotter_circuit=PauliTrotterEvolution().convert(time_evolution_operator)  #trotterize it.
bound = trotter_circuit.bind_parameters({t: time_step_value_t}) #bind it
#tight_binding_circuit.append(bound,[[0],[1],[2]])
for _ in range(steps_t):
    tight_binding_circuit.append(bound, [[0], [1], [2]])
full_time_evolution_circuit=tight_binding_circuit

###DO NOT EDIT BELOW THIS LINE

print(full_time_evolution_circuit)



full_time_evolution_circuit.decompose().decompose().decompose().draw() # create a trotterized decomposition

from qc_grader.challenges.qgss_2022 import grade_lab2_ex12

grade_lab2_ex12(full_time_evolution_circuit)