import numpy as np
import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 16})  # enlarge matplotlib fonts

# Import qubit states Zero (|0>) and One (|1>), Pauli operators (X, Y, Z), and the identity operator (I)
from qiskit.opflow import Zero, One, X, Y, Z, I

# Returns the XXX Heisenberg model for 3 spin-1/2 particles in a line
def ex1_compute_H_heis3():
    # FILL YOUR CODE IN HERE
    H = (I^X^X) + (X^X^I) + (I^Y^Y) + (Y^Y^I) + (I^Z^Z) + (Z^Z^I)
    # Return Hamiltonian
    return H

from qc_grader.challenges.qgss_2022 import grade_lab4_ex1

# The grading function is expecting a PauliSumOp operator
grade_lab4_ex1(ex1_compute_H_heis3())

# Returns the time evolution operator U_heis3(t) for a given time t assuming an XXX Heisenberg Hamiltonian for 3 spins-1/2 particles in a line
def ex2_compute_U_heis3(t):
    # FILL YOUR CODE IN HERE
    H = ex1_compute_H_heis3()
    U = (t*H).exp_i()
    return U

from qc_grader.challenges.qgss_2022 import grade_lab4_ex2

# The grading function is expecting a funtion that returns PauliSumOp operator
grade_lab4_ex2(ex2_compute_U_heis3)

# Define array of time points
ts = np.linspace(0, np.pi, 100)

# Define initial state |110>
initial_state = One^One^Zero

# Compute probability of remaining in |110> state over the array of time points
 # ~initial_state gives the bra of the initial state (<110|)
 # @ is short hand for matrix multiplication
 # ex2_compute_U_heis3(t) is the unitary time evolution at time t
 # t needs to be wrapped with float(t) to avoid a bug
 # (...).eval() returns the inner product <110|ex2_compute_U_heis3(t)|110>
 #  np.abs(...)**2 is the modulus squared of the innner product which is the expectation value, or probability, of remaining in |110>
probs_110 = [np.abs((~initial_state @ ex2_compute_U_heis3(float(t)) @ initial_state).eval())**2 for t in ts]

# Plot evolution of |110>
plt.plot(ts, probs_110, linewidth=2)
plt.xlabel('time')
plt.ylabel(r'probability of state $|110\rangle$')
plt.title(r'Evolution of state $|110\rangle$ under $H_{Heis3}$')
plt.ylim([-0.05,1.05])
plt.grid()
plt.show()

!pip install qiskit_ignis --quiet

# Importing standard Qiskit modules
from qiskit import QuantumCircuit, QuantumRegister, IBMQ, execute, transpile
from qiskit.providers.aer import QasmSimulator
from qiskit.tools.monitor import job_monitor
from qiskit.circuit import Parameter

# Import state tomography modules
from qiskit.ignis.verification.tomography import state_tomography_circuits, StateTomographyFitter
from qiskit.quantum_info import state_fidelity

# load IBMQ Account data

# IBMQ.save_account(TOKEN)  # replace TOKEN with your API token string (https://quantum-computing.ibm.com/lab/docs/iql/manage/account/ibmq)
provider = IBMQ.load_account()

# Get backend for experiment
provider = IBMQ.get_provider(hub='ibm-q')
manila = provider.get_backend('ibmq_manila')
# properties = manila.properties()

# Simulated backend based on ibmq_manila's device noise profile
sim_noisy_manila = QasmSimulator.from_backend(provider.get_backend('ibmq_manila'))

# Noiseless simulated backend
sim = QasmSimulator()

# Parameterize variable t to be evaluated at t=pi later
t = Parameter('t')

# Build a subcircuit for XX(t) two-qubit gate
def compute_XX_gate(t):
    XX_qr = QuantumRegister(2)
    XX_qc = QuantumCircuit(XX_qr, name='XX')

    XX_qc.ry(np.pi/2,[0,1])
    XX_qc.cnot(0,1)
    XX_qc.rz(2 * t, 1)
    XX_qc.cnot(0,1)
    XX_qc.ry(-np.pi/2,[0,1])

    # Convert custom quantum circuit into a gate
    XX = XX_qc.to_instruction()
    
    return XX

# Build a subcircuit for YY(t) two-qubit gate
def ex3_compute_YY_gate(t):
    # FILL YOUR CODE IN HERE

    YY_qr = QuantumRegister(2)
    YY_qc = QuantumCircuit(YY_qr, name='YY')
    YY_qc.s(0)
    YY_qc.s(1)
    XX = compute_XX_gate(t)
    YY_qc.append(XX, [YY_qr[0], YY_qr[1]])
    YY_qc.sdg(0)
    YY_qc.sdg(1)
    
    YY = YY_qc.to_instruction()
    
    return YY

from qc_grader.challenges.qgss_2022 import grade_lab4_ex3

# The grading function is expecting an Instruction
grade_lab4_ex3(ex3_compute_YY_gate(t))

# Build a subcircuit for ZZ(t) two-qubit gate
def ex4_compute_ZZ_gate(t):
    # FILL YOUR CODE IN HERE
    
    ZZ_qr = QuantumRegister(2)
    ZZ_qc = QuantumCircuit(ZZ_qr, name='ZZ')
    ZZ_qc.h(0)
    ZZ_qc.h(1)
    XX = compute_XX_gate(t)
    ZZ_qc.append(XX, [ZZ_qr[0], ZZ_qr[1]])
    ZZ_qc.h(0)
    ZZ_qc.h(1)
    
    ZZ = ZZ_qc.to_instruction()
    
    return ZZ

from qc_grader.challenges.qgss_2022 import grade_lab4_ex4

# The grading function is expecting an Instruction
grade_lab4_ex4(ex4_compute_ZZ_gate(t))

# Combine subcircuits into a single multiqubit gate representing a single trotter step
num_qubits = 3

# Define two-qubit interactions Parameterized by t
XX = compute_XX_gate(t)
YY = ex3_compute_YY_gate(t)
ZZ = ex4_compute_ZZ_gate(t)

Trot_qr = QuantumRegister(num_qubits)
Trot_qc = QuantumCircuit(Trot_qr, name='Trot')

for i in range(0, num_qubits - 1):
    Trot_qc.append(ZZ, [Trot_qr[i], Trot_qr[i+1]])
    Trot_qc.append(YY, [Trot_qr[i], Trot_qr[i+1]])
    Trot_qc.append(XX, [Trot_qr[i], Trot_qr[i+1]])

# Convert custom quantum circuit into a gate
Trot_gate = Trot_qc.to_instruction()

# Setup experiment parameters

# The final time of the state evolution
target_time = np.pi  # DO NOT MODIFY

# Number of trotter steps
trotter_steps = 20  ### CAN BE >= 4

# Select which qubits to use for the simulation
q_regs = [0,1,2]

# Initialize quantum circuit for 3 qubits
qr = QuantumRegister(5)
qc = QuantumCircuit(qr)

# Prepare initial state (remember we are only evolving 3 of the 7 qubits on manila qubits (q_5, q_3, q_1) corresponding to the state |110>)
qc.x([q_regs[2], q_regs[1]])  # For example this could be (q_regs=[2, 1, 0] which corresponds to => |110>)

# Simulate time evolution under H_heis3 Hamiltonian
for _ in range(trotter_steps):
    qc.append(Trot_gate, q_regs)

# Evaluate simulation at target_time (t=pi) meaning each trotter step evolves pi/trotter_steps in time
qc = qc.bind_parameters({t: target_time/trotter_steps})

# Generate state tomography circuits to evaluate fidelity of simulation
st_qcs = state_tomography_circuits(qc, q_regs)

# Display circuit for confirmation
# st_qcs[-1].decompose().draw()  # view decomposition of trotter gates
st_qcs[-1].draw()  # only view trotter gates

shots = 8192
reps = 4

# Pick the simulated or real backend for manila
# backend = sim  # Noiseless simulator
backend = sim_noisy_manila  # Simulator that mimics ibmq_manila
# backend = manila  # The real manila backend

jobs = []
for _ in range(reps):
    # Execute
    job = execute(st_qcs, backend, shots=shots)
    print('Job ID', job.job_id())
    jobs.append(job)

for job in jobs:
    job_monitor(job)
    try:
        if job.error_message() is not None:
            print(job.error_message())
    except:
        pass

# Compute the state tomography based on the st_qcs quantum circuits and the results from those ciricuits
def state_tomo(result, st_qcs):
    # The expected final state; necessary to determine state tomography fidelity
    target_state = (One^One^Zero).to_matrix()  # DO NOT MODIFY
    # Fit state tomography results
    tomo_fitter = StateTomographyFitter(result, st_qcs)
    rho_fit = tomo_fitter.fit(method='lstsq')
    # Compute fidelity
    fid = state_fidelity(rho_fit, target_state)
    return fid

# Compute tomography fidelities for each repetition
fids = []
for job in jobs:
    fid = state_tomo(job.result(), st_qcs)
    fids.append(fid)
    
# Share tomography fidelity of discord to compete and collaborate with other students
print('state tomography fidelity on ' + str(backend) + ' = {:.4f} \u00B1 {:.4f}'.format(np.mean(fids), np.std(fids)))
# Share what lectures and techniques were useful in optimizing your results
print('Inspiration: Measurement error mitigation, Olivia Lanes\'s 2nd lecture')

import qiskit.tools.jupyter
%qiskit_version_table

