import numpy as np
from scipy.linalg import expm

from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit, Aer, execute

from qiskit.circuit import AncillaQubit

from qiskit.visualization import plot_histogram

from qiskit.circuit.library.standard_gates import XGate, ZGate, HGate

from qiskit.circuit.add_control import add_control

from qiskit.extensions import UnitaryGate

def test_circuit(qc):
    ''' 
    auxiliar function, to test intermediate steps, looking at measurement results
    this function allows the inspection of the amplitudes of states, 
    at any point of the circuit (call it with the desired circuit)
    '''
    
    cr = ClassicalRegister(qc.num_qubits, "m_cl")
    qc_cr = QuantumCircuit(cr)

    qc_test = qc + qc_cr

    for i in range(qc.num_qubits):
        qc_test.measure(i, i)

    #################################

    backend = Aer.get_backend("qasm_simulator")

    job = execute(qc_test, backend, shots=1e5, seed_simulator=42)
    results = job.result()
    counts = results.get_counts()

    return plot_histogram(counts, title="Results", figsize=(12, 4))

qr_input = QuantumRegister(4, "i")
qc_input = QuantumCircuit(qr_input)

x = "0010"
for i in range(len(x)):
    if x[i] == "1":
        qc_input.x(i)
        
qc_input.draw("mpl")

test_circuit(qc_input)

X = np.array([[0, 0, 0, 0],
              [0, 0, 0, 1],
              [1, 1, 1, 0],
              [1, 1, 1, 1]])

# number of examples
N = X.shape[0]

# number of features
n = X.shape[1]

y = np.array([0, 0, 1, 1])

# full dataset, of the form [[X], [y]]
dataset = np.append(X, y.reshape((-1, 1)), axis=1)

# n+1 because the last register will encode the class!

amplitudes = np.zeros(2**(n+1))

amplitudes

# integer representation of binary strings in training dataset
# notice the [::-1], which is necessary to adjust the endianness!

data_points_X = [int("".join(str(i) for i in X[j])[::-1], 2) for j in range(dataset.shape[0])]

data_points_X

# integer representation considering also the class
# IMPORTANT: sum 2**n for elements in class 1
# this is necessary to encode also the class!

data_points = [x + 2**n if y[i] == 1 else x for i, x in enumerate(data_points_X)]

data_points

amplitudes[data_points] = 1

amplitudes = amplitudes/np.sqrt(amplitudes.sum())

amplitudes

qr_features = QuantumRegister(n, "f")
qr_class = QuantumRegister(1, "cl")

qc_training = QuantumCircuit(qr_features, qr_class)

qc_training.initialize(amplitudes, [qr_features, qr_class])

qc_training.draw("mpl")

test_circuit(qc_training)

dataset

qr_anc = QuantumRegister(1, "anc")
qc_anc = QuantumCircuit(qr_anc)

qc_anc.draw("mpl")

qc = qc_input + qc_training + qc_anc

qc.barrier()

qc.draw("mpl")

test_circuit(qc)

qc.h(-1)

qc.barrier()

qc.draw("mpl")

for i in range(4):
    
    qc.cnot(i, 4+i)
    
qc.barrier()
    
qc.draw("mpl")

# matrices of 1 and \sigma_z to be exponentiated below
# source of the matrix exponential methods below:
# https://quantumcomputing.stackexchange.com/questions/10317/quantum-circuit-to-implement-matrix-exponential

idtt = np.eye(2)

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

# define the exponentiated unitaries
U_1_minus = expm(-1j * (np.pi/(2*4)) * ((idtt-sz)/2))
U_1_plus = expm(1j * (np.pi/(2*4)) * ((idtt-sz)/2))

# defining controlled gates
u1m = add_control(operation=UnitaryGate(U_1_minus, label="$U_1^-$"),
                  num_ctrl_qubits=1, ctrl_state=0, label="$CU_1^-$")

u1p = add_control(operation=UnitaryGate(U_1_plus, label="$U_1^+$"),
                  num_ctrl_qubits=1, ctrl_state=1, label="$CU_1^+$")


# build a circuit to apply the unitary above
# this will be combined with the main circuit later on (notice: same registers!!).
qc_u = QuantumCircuit(qr_features, qr_anc)
           
for i in range(4):
    
    # apply the unitaries
    qc_u.append(u1p, [qr_anc[0], qr_features[i]])
    qc_u.append(u1m, [qr_anc[0], qr_features[i]])
      
# # decompose this portion of the circuit
# # (this will be appended to the main circuit later on)
# qc_u = qc_u.decompose()

qc_u.draw("mpl")

qr_input = QuantumRegister(4, "i")
qr_features = QuantumRegister(4, "f")
qr_class = QuantumRegister(1, "cl")
qr_anc = QuantumRegister(1, "anc")

qc_prep = QuantumCircuit(qr_input, qr_features, qr_class, qr_anc)

# build the state |0010,0110,0 ;1⟩ + |0010,0110,0 ;0⟩

# input
qc_prep.x(qr_input[2])

# hamming ditance
# qc_prep.x(qr_features[1])
# qc_prep.x(qr_features[2])
qc_prep.x(qr_features[3])

# ancilla
qc_prep.h(qr_anc[0])

qc_prep.barrier()

##################################################

qc_test = qc_prep.combine(qc_u)

qc_test.barrier()

##################################################

qc_test.h(qr_anc[0])

qc_test.barrier()

##################################################

cr = ClassicalRegister(qc_test.num_qubits, "c")
qc_cr = QuantumCircuit(cr)

qc_test = qc_test + qc_cr

for i in range(qc_test.num_qubits):
    
    qc_test.measure(i, i)
    
##################################################

qc_test.draw("mpl")

cr = ClassicalRegister(qc_test.num_qubits, "c")
qc_cr = QuantumCircuit(cr)

qc_prep = qc_prep + qc_cr

for i in range(qc_prep.num_qubits):
    
    qc_prep.measure(i, i)

backend = Aer.get_backend("qasm_simulator")

job = execute(qc_prep, backend, shots=1e5, seed_simulator=42)
results = job.result()
counts = results.get_counts()

plot_histogram(counts, title="Results", figsize=(12, 4))

backend = Aer.get_backend("qasm_simulator")

job = execute(qc_test, backend, shots=1e5, seed_simulator=42)
results = job.result()
counts = results.get_counts()

plot_histogram(counts, title="Results", figsize=(12, 4))

print("Ancila 0")

(np.cos(np.pi/8))**2

print("Ancila 1")

(np.sin(np.pi/8))**2

qc = qc.combine(qc_u)

qc.barrier()

qc.draw("mpl")

qc.h(-1)

qc.barrier()

qc.draw("mpl", scale=0.75)

# copying the qc so far, before measurement
qc_bckp = qc.copy()

qc = qc_bckp.copy()

cr = ClassicalRegister(qc.num_qubits, "c")
qc_cr = QuantumCircuit(cr)

qc = qc + qc_cr

for i in range(qc.num_qubits):
    
    qc.measure(i, i)
    
# qc.draw("mpl")

backend = Aer.get_backend("qasm_simulator")

job = execute(qc, backend, shots=1e5, seed_simulator=42)
results = job.result()
counts = results.get_counts()

plot_histogram(counts, title="Results", figsize=(12, 4))

# using the same endiannes of qiskit, i.e., |0010,0010,0;0> = 0001000100

{"0001000100" : round((0.25)*(np.cos(np.pi/8))**2, 3),
 "1001000100" : round((0.25)*(np.sin(np.pi/8))**2, 3),
 
 "0011000100" : round((0.25)*(np.cos(2*np.pi/8))**2, 3),
 "1011000100" : round((0.25)*(np.sin(2*np.pi/8))**2, 3),
 
 "0100110100" : round((0.25)*(np.cos(2*np.pi/8))**2, 3),
 "1100110100" : round((0.25)*(np.sin(2*np.pi/8))**2, 3),
 
 "0110110100" : round((0.25)*(np.cos(3*np.pi/8))**2, 3),
 "1110110100" : round((0.25)*(np.sin(3*np.pi/8))**2, 3)}

qc = qc_bckp.copy()

# first, lets's measure the ancilla

cr_anc = ClassicalRegister(1, "c_anc")
qc_cr_anc = QuantumCircuit(cr_anc)

qc_cr_anc.draw("mpl")

qc = qc + qc_cr_anc

qc.measure(qr_anc, cr_anc[0])

# qc.draw("mpl", scale=0.7)

# theoretical probability of measuring ancilla in |0>

(0.25)*((np.cos(np.pi/8))**2 + 2*(np.cos(2*np.pi/8))**2 + (np.cos(3*np.pi/8))**2)

backend = Aer.get_backend("qasm_simulator")

job = execute(qc, backend, shots=1e4, seed_simulator=42)
results = job.result()
counts = results.get_counts()

plot_histogram(counts, title="Results", figsize=(12, 4))

qc = qc_bckp.copy()

cr_anc = ClassicalRegister(1, "c_anc")
qc_cr_anc = QuantumCircuit(cr_anc)

cr_class = ClassicalRegister(1, "c_class")
qc_cr_class = QuantumCircuit(cr_class)

qc = qc + qc_cr_anc + qc_cr_class

qc.measure(qr_anc, cr_anc)
qc.measure(qr_class, cr_class)

qc.draw("mpl")

backend = Aer.get_backend("qasm_simulator")

job = execute(qc, backend, shots=1e4, seed_simulator=42)
results = job.result()
counts = results.get_counts()

# filtering out measurements of ancila = |1> (last register)
keys = list(counts.keys()).copy()
for key in keys:
    
    if key[-1] == "1":
        
        counts.pop(key)

plot_histogram(counts, title="Results", figsize=(12, 4))

# theoretical probability of measuring class qubit in |0>, given ancilla was in |0>

(0.5)*((np.cos(np.pi/8))**2 + (np.cos(2*np.pi/8))**2)

# probability of measuring class qubit in |1>, given ancilla was in |0>

(0.5)*((np.cos(2*np.pi/8))**2 + (np.cos(3*np.pi/8))**2)