"""Imports for the notebook."""
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import minimize

import qiskit

"""Defining circuits with parameters."""
# Get a circuit and registers
qreg = qiskit.QuantumRegister(2)
creg = qiskit.ClassicalRegister(2)
circ = qiskit.QuantumCircuit(qreg, creg)

# Add gates with particular parameters
circ.h(qreg)
circ.rx(0.2, qreg[0])
circ.cx(qreg[0], qreg[1])
circ.ry(0.1, qreg[1])

# Visualize the circuit
print(circ)

def circuit(alpha1: float, alpha2: float):
    """Returns the circuit above with the input parameters."""
    ### Your code here!

"""Estimating a one qubit expectation value."""
qreg = qiskit.QuantumRegister(1)
creg = qiskit.ClassicalRegister(1)
circ = qiskit.QuantumCircuit(qreg, creg)
circ.h(qreg)

### Your code here!

"""Estimating a one qubit expectation value."""
qreg = qiskit.QuantumRegister(1)
creg = qiskit.ClassicalRegister(1)
circ = qiskit.QuantumCircuit(qreg, creg)
circ.x(qreg)
circ.h(qreg)

### Your code here!

"""Estimating a two qubit expectation value."""
qreg = qiskit.QuantumRegister(2)
creg = qiskit.ClassicalRegister(2)
circ = qiskit.QuantumCircuit(qreg, creg)
circ.h(qreg[1])

### Your code here!

"""Helper function to evaluate the expectation of any valid Pauli string."""
def expectation_circuit(circuit: qiskit.QuantumCircuit, pauli_string: str) -> qiskit.QuantumCircuit:
    """Returns a circuit to compute expectation of the Pauli string in the 
    state prepared by the input circuit.
    
    Args:
        circuit: Prepares the state |\psi> from |0>.
        pauli_string: String (tensor product) of Paulis to evaluate
                      an expectation of. The length of pauli_string
                      must be equal to the total number of qubits in
                      the circuit. (Use identities for no operator!)
    """
    if len(circuit.qregs) != 1:
        raise ValueError("Circuit should have only one quantum register.")
    if len(circuit.cregs) != 1:
        print("# cregs =", len(circuit.cregs))
        raise ValueError("Circuit should have only one classical register.")
    ### Your code here!

"""Test your function here."""
circ, qreg, creg = circuit(np.pi / 2, np.pi / 4)
print("Bare circuit:")
print(circ)
### Your code here!

"""Function to execute the circuit and do the postprocessing."""
def run_and_process(circuit: qiskit.QuantumCircuit, shots: int = 10000) -> float:
    """Runs an 'expectation circuit' and returns the expectation value of the
    measured Pauli string.
    
    Args:
        circuit: Circuit to execute.
        shots: Number of circuit executions.
    """
    ### Your code here!

"""Define your function here!"""
def expectation(circuit: qiskit.QuantumCircuit, pauli_string: str, shots: int = 10000) -> float:
    """Returns the expectation value of the pauli string in the state prepared by the circuit."""
    ### Your code here!

"""Test your function here."""
### Your code here!

"""Compute the expectation of a Hamiltonian."""
# Provided circuit
qreg = qiskit.QuantumRegister(3)
creg = qiskit.ClassicalRegister(3)
circ = qiskit.QuantumCircuit(qreg, creg)
circ.h(qreg)
circ.rx(np.pi / 4, qreg[0])
circ.cz(qreg[0], qreg[1])
circ.cz(qreg[1], qreg[2])
print(circ)

weights = (0.5, -0.3, 1.2)
paulis = ("IZZ", "ZZI", "ZIZ")

### Your code here

"""Function to compute the cost of any Hamiltonian in the state prepared by the circuit."""
def cost(circuit, weights, paulis):
    """Returns <psi|H|psi> where |psi> is prepared by the circuit
    and the weights and paulis define a Hamiltonian H.
    
    Args:
        circuit: Circuit which prepares a state.
        weights: List of floats which are the coeffs/weights of each Pauli string.
        paulis:  List of strings which specify the Paulis.
    """
    if len(weights) != len(paulis):
        raise ValueError("Args weights and paulis must have the same length.")
    ### Your code here!

"""Evaluate your cost here!"""
### Your code here!

"""Plot a cost landscape."""
def oneq_circ(param):
    qreg = qiskit.QuantumRegister(1)
    creg = qiskit.ClassicalRegister(1)
    circ = qiskit.QuantumCircuit(qreg, creg)
    circ.rx(param, qreg)
    return circ

weights = (1.0,)
paulis = ("Z",)

pvals = np.linspace(-np.pi, np.pi, 100)
cvals = []

### Your code here!

"""Get a parameterized circuit."""
def circuit(params):
    qreg = qiskit.QuantumRegister(2)
    creg = qiskit.ClassicalRegister(2)
    circ = qiskit.QuantumCircuit(qreg, creg)
    circ.h(qreg)
    circ.cx(qreg[0], qreg[1])
    circ.rx(params[0], qreg[0])
    circ.ry(params[1], qreg[1])
    circ.cz(qreg[0], qreg[1])
    circ.s(qreg[0])
    circ.t(qreg[1])
    return circ

"""Visualize the circuit."""
print(circuit([1, 2]))

"""Hamiltonian cost to minimize."""
weights = (1.2, -0.2)
paulis = ("IZ", "ZX")

"""Define a cost/objective function."""
def obj(params):
    """Returns the cost for the given parameters."""
    ### Your code here

"""Test your function on this set of parameters."""
obj([0, 0])

"""Run an optimization algorithm to return the lowest cost and best parameters."""
result = minimize(obj, x0=[0, 0], method="COBYLA")

"""See the optimization results."""
print("Lowest cost function value found:", result.fun)
print("Best parameters:", result.x)