import numpy as np
from numpy import kron
import qiskit as qk

def decomposition(H):
    """Decompose any 4x4 Hermitian matrix into a sum of tensor products of two Pauli matrices
    """
    identity = np.array([[1, 0],[ 0, 1]], dtype=np.complex128)
    x = np.array([[0, 1],  [ 1, 0]], dtype=np.complex128)
    y = np.array([[0, -1j],[1j, 0]], dtype=np.complex128)
    z = np.array([[1, 0],  [0, -1]], dtype=np.complex128)
    S = [identity, x, y, z]
    labels = ['I', 'sigma_x', 'sigma_y', 'sigma_z']
    d = {'00': 'I product I', '01': 'I product sigma_x', '02': 'I product sigma_y', '03': 'I product sigma_z', 
        '10': 'sigma_x product I', '11': 'sigma_x product sigma_x', '12': 'sigma_x product sigma_y', 
        '13': 'sigma_x product sigma_z', '20': 'sigma_y product I', '21': 'sigma_y product sigma_x', 
        '22': 'sigma_y product sigma_y', '23': 'sigma_y product sigma_z', '30': 'sigma_z product I', 
        '31': 'sigma_z product sigma_x', '32': 'sigma_z product sigma_y', '33': 'sigma_z product sigma_z'}
    for i in range(4):
        for j in range(4):
            a_ij = 0.25 * np.dot(kron(S[i], S[j]), H).trace()
            if a_ij != 0.0:
                print(str(a_ij) + ' * ' + d[str(i)+str(j)])

H = np.array([[0,0,0,0],[0,-1,1,0],[0,1,-1,0],[0,0,0,0]])
decomposition(H)

def prepare_state(theta, n=3):
    """
    Prepare three 2-qubit states with 3 associated quantum registers, 3 associated classical registers, 
    and 3 quantum circuits.  We will prepare the state with the ansatz mentioned in the notes in which we act on
    the first qubit with the Hadamard operator, then with the R_z operator, then we act on the 2-qubit state
    with the CNOT gate, and then on the second qubit in each terms of the superposition with the sigma_x operator.
    After that, we will take the first circuit and act on each qubit with the R_y(pi/2) operator, and take the second
    circuit and act on each qubit with the R_x(-pi/2) operator as explained in the notes. We do this so that those
    qubits will be in the basis of eigenvectors of sigma_x and sigma_y as explained in the notes.  We can measure
    the qubits in the other circuit as-is because we need the expectation value of sigma_z and the qubits are already
    in the computational basis.
    """
    qr0 = qk.QuantumRegister(2)
    cr0 = qk.ClassicalRegister(2)
    qc0 = qk.QuantumCircuit(qr0,cr0)
    qr1 = qk.QuantumRegister(2)
    cr1 = qk.ClassicalRegister(2)
    qc1 = qk.QuantumCircuit(qr1,cr1)
    qr2 = qk.QuantumRegister(2)
    cr2 = qk.ClassicalRegister(2)
    qc2 = qk.QuantumCircuit(qr2,cr2)
    
    qregisters = [qr0,qr1,qr2]
    cregisters = [cr0,cr1,cr2]
    qcircuits = [qc0,qc1,qc2]
    
    for i in range(n):
        qcircuits[i].h(qregisters[i][0])
    for i in range(n):
        qcircuits[i].rz(theta, qregisters[i][0])
    for i in range(n):
        qcircuits[i].cx(qregisters[i][0], qregisters[i][1])
    for i in range(n):
        qcircuits[i].x(qregisters[i][1])
        
    qcircuits[0].ry((np.pi)/2, qregisters[0][0])
    qcircuits[0].ry((np.pi)/2, qregisters[0][1])
    qcircuits[1].rx(-(np.pi)/2, qregisters[1][0])
    qcircuits[1].rx(-(np.pi)/2, qregisters[1][1])
    
    return qregisters, cregisters, qcircuits

qregisters, cregisters, qcircuits = prepare_state(np.pi, n=3)
qcircuits[0].draw(output='mpl')

qcircuits[1].draw(output='mpl')

qcircuits[2].draw(output='mpl')

def expectation(qcircuits, cregisters, qregisters, n_shots, n=3):
    """
    For each circuit, execute and measure it using the classical simulator 5000 times as explained above. Multiply
    each of the three expectation values by 0.5, add them up, and subtract -0.5.  Return the result.
    """
    expect = -0.5
    for i in range(n):
        qcircuits[i].measure(qregisters[i],cregisters[i])
        qk.Aer.backends()
        sim = qk.Aer.get_backend('qasm_simulator')
        res = qk.execute(qcircuits[i], sim, shots=n_shots).result()
        counts = res.get_counts()
        sum = 0
        for k,v in counts.items():
            if k=='01' or k=='10':
                sum += (-1)*v/n_shots
            elif k=='00' or k=='11':
                sum += v/n_shots
        sum = 0.5*sum
        expect += sum
    return expect

# Consider 100 values of theta, between 0 and Pi.  This theta is the one used in state preparation.
thetas = np.linspace(0, np.pi, 100)

# For each theta, store the resulting expectation value in results
results = []
# For each theta, find the expectation value
for theta in thetas:
    qregisters, cregisters, qcircuits = prepare_state(theta, n=3)
    expect = expectation(qcircuits, cregisters, qregisters, 5000, n=3)
    results.append(expect)
# Sort the results in ascending order.  The first one is your minimum eigenvalue.
results.sort()
print("The minimum eigenvalue is: {}.".format(results[0]))

from qiskit import IBMQ
#IBMQ.delete_account()
IBMQ.save_account('my IBM token', overwrite=True)
IBMQ.load_account()

provider = IBMQ.get_provider()
procs=provider.backends(operational=True, simulator=False)

from qiskit.tools.jupyter import *
%qiskit_backend_overview

from qiskit.tools import monitor
backend = qk.providers.ibmq.least_busy([p for p in procs if len(p.properties().qubits) >= 2])

from qiskit.tools.monitor import backend_overview, backend_monitor
backend_monitor(backend)

def q_expectation(qcircuits, cregisters, qregisters, n_shots, n=3):
    """
    For each circuit, execute and measure it using the classical simulator 5000 times as explained above. Multiply
    each of the three expectation values by 0.5, add them up, and subtract -0.5.  Return the result.
    """
    expect = -0.5
    for i in range(n):
        qcircuits[i].measure(qregisters[i],cregisters[i])
        qk.Aer.backends()
        sim = qk.Aer.get_backend('qasm_simulator')
        res = qk.execute(qcircuits[i], backend=backend, shots=n_shots).result()
        #mon = monitor.job_monitor(res)
        counts = res.get_counts()
        sum = 0
        for k,v in counts.items():
            if k=='01' or k=='10':
                sum += (-1)*v/n_shots
            elif k=='00' or k=='11':
                sum += v/n_shots
        sum = 0.5*sum
        expect += sum
    return expect

# Consider 10 values of theta, between 0 and Pi.  This theta is the one used in state preparation.
thetas = np.linspace(0, np.pi, 10)

# Use n_shots = 100
results = []
# For each theta, find the expectation value
for theta in thetas:
    qregisters, cregisters, qcircuits = prepare_state(theta, n=3)
    expect = q_expectation(qcircuits, cregisters, qregisters, 100, n=3)
    results.append(expect)
# Sort the results in ascending order.  The first one is your minimum eigenvalue.
results.sort()
print("The minimum eigenvalue is: {}.".format(results[0]))

# Use n_shots = 1000
results = []
# For each theta, find the expectation value
for theta in thetas:
    qregisters, cregisters, qcircuits = prepare_state(theta, n=3)
    expect = q_expectation(qcircuits, cregisters, qregisters, 1000, n=3)
    results.append(expect)
# Sort the results in ascending order.  The first one is your minimum eigenvalue.
results.sort()
print("The minimum eigenvalue is: {}.".format(results[0]))

# Use n_shots = 5000
results = []
# For each theta, find the expectation value
for theta in thetas:
    qregisters, cregisters, qcircuits = prepare_state(theta, n=3)
    expect = q_expectation(qcircuits, cregisters, qregisters, 5000, n=3)
    results.append(expect)
# Sort the results in ascending order.  The first one is your minimum eigenvalue.
results.sort()
print("The minimum eigenvalue is: {}.".format(results[0]))

# Use n_shots = 8192 = max allowed
results = []
thetas = [np.pi]
# For each theta, find the expectation value
for theta in thetas:
    qregisters, cregisters, qcircuits = prepare_state(theta, n=3)
    expect = q_expectation(qcircuits, cregisters, qregisters, 8192, n=3)
    results.append(expect)
# Sort the results in ascending order.  The first one is your minimum eigenvalue.
results.sort()
print("The minimum eigenvalue is: {}.".format(results[0]))

def decomposition(H):
    """Decompose any 4x4 Hermitian matrix into a sum of tensor products of two Pauli matrices
    """
    A_ij = []
    identity = np.array([[1, 0],[ 0, 1]], dtype=np.complex128)
    x = np.array([[0, 1],  [ 1, 0]], dtype=np.complex128)
    y = np.array([[0, -1j],[1j, 0]], dtype=np.complex128)
    z = np.array([[1, 0],  [0, -1]], dtype=np.complex128)
    S = [identity, x, y, z]
    labels = ['I', 'sigma_x', 'sigma_y', 'sigma_z']
    d = {'00': 'I product I', '01': 'I product sigma_x', '02': 'I product sigma_y', '03': 'I product sigma_z', 
        '10': 'sigma_x product I', '11': 'sigma_x product sigma_x', '12': 'sigma_x product sigma_y', '13': 'sigma_x product sigma_z',
        '20': 'sigma_y product I', '21': 'sigma_y product sigma_x', '22': 'sigma_y product sigma_y', '23': 'sigma_y product sigma_z',
        '30': 'sigma_z product I', '31': 'sigma_z product sigma_x', '32': 'sigma_z product sigma_y', '33': 'sigma_z product sigma_z'}
    for i in range(4):
        for j in range(4):
            a_ij = 0.25 * np.dot(kron(S[i], S[j]), H).trace()
            A_ij.append(a_ij)
            if a_ij != 0.0:
                print(str(a_ij) + ' * ' + d[str(i)+str(j)])
    return np.asarray(A_ij).reshape(4,4)

def prepare_state2(A, theta):
    qregisters = []
    cregisters = []
    qcircuits = []
    identity = np.array([[1, 0],[ 0, 1]], dtype=np.complex128)
    x = np.array([[0, 1],  [ 1, 0]], dtype=np.complex128)
    y = np.array([[0, -1j],[1j, 0]], dtype=np.complex128)
    z = np.array([[1, 0],  [0, -1]], dtype=np.complex128)
    d = {}
    for i in range(4):
        for j in range(4):
            if A[i,j] != 0:
                if i !=0 and j!=0:
                    qr = qk.QuantumRegister(2)
                    cr = qk.ClassicalRegister(2)
                    qc = qk.QuantumCircuit(qr,cr)
                    qc.h(qr[0])
                    qc.rz(theta, qr[0])
                    qc.cx(qr[0], qr[1])
                    qc.x(qr[1])
                    if i==1:
                        qc.ry((np.pi)/2, qr[0])
                    if i==2:
                        qc.rx(-(np.pi)/2, qr[0])
                    if j==1:
                        qc.ry((np.pi)/2, qr[1])
                    if j==2:
                        qc.rx(-(np.pi)/2, qr[1])
                    qregisters.append(qr)
                    cregisters.append(cr)
                    qcircuits.append(qc)
                    d[(i,j)] = [qregisters, cregisters, qcircuits]
                if i == 0 and j != 0:
                    qr = qk.QuantumRegister(2)
                    cr = qk.ClassicalRegister(2)
                    qc = qk.QuantumCircuit(qr,cr)
                    qc.h(qr[0])
                    qc.rz(theta, qr[0])
                    qc.cx(qr[0], qr[1])
                    qc.x(qr[1])
                    if j==1:
                        qc.ry((np.pi)/2, qr[1])
                    if j==2:
                        qc.rx(-(np.pi)/2, qr[1])
                    qregisters.append(qr)
                    cregisters.append(cr)
                    qcircuits.append(qc)
                    d[(i,j)] = [qregisters, cregisters, qcircuits]
                if i != 0 and j == 0:
                    qr = qk.QuantumRegister(2)
                    cr = qk.ClassicalRegister(2)
                    qc = qk.QuantumCircuit(qr,cr)
                    qc.h(qr[0])
                    qc.rz(theta, qr[0])
                    qc.cx(qr[0], qr[1])
                    qc.x(qr[1])
                    if i==1:
                        qc.ry((np.pi)/2, qr[0])
                    if i==2:
                        qc.rx(-(np.pi)/2, qr[0])
                    qregisters.append(qr)
                    cregisters.append(cr)
                    qcircuits.append(qc)
                    d[(i,j)] = [qregisters, cregisters, qcircuits]
    return d, A

def expectation2(d, A, n_shots):
    """
    For each circuit, execute and measure it using the classical simulator 5000 times as explained above. 
    Return the result.
    """
    expect = A[0,0]
    qk.Aer.backends()
    sim = qk.Aer.get_backend('qasm_simulator')
    for k,v in d.items():
        for i in range(len(v[2])):
            v[2][i].measure(v[0][i],v[1][i])
            res = qk.execute(v[2][i], sim, shots=n_shots).result()
            counts = res.get_counts()
            sum = 0
            for m,n in counts.items():
                if m=='01' or m=='10':
                    sum += (-1)*n/n_shots
                elif m=='00' or m=='11':
                    sum += n/n_shots
            sum = A[k[0],k[1]]*sum
            expect += sum
    return expect

thetas = np.linspace(0, np.pi, 100)
H = np.array([[0,0,0,0],[0,-1,1,0],[0,1,-1,0],[0,0,0,0]])
A = decomposition(H)

# For each theta, store the resulting expectation value in results
results = []
# For each theta, find the expectation value
for theta in thetas:
    d, A = prepare_state2(A, theta)
    expect = expectation2(d, A, 5000)
    results.append(expect)
# Sort the results in ascending order.  The first one is your minimum eigenvalue.
results.sort()
print("The minimum eigenvalue is: {}.".format(results[0]))

# -5??!!  Maybe I missed something in my logic.  Will continue to look into this. :-)