import numpy as np
from qiskit import *
from qiskit.tools.jupyter import *
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from sklearn.preprocessing import Normalizer

backend = BasicAer.get_backend('qasm_simulator')


def get_angles(x):

    beta0 = 2 * np.arcsin(np.sqrt(x[1]) ** 2 / np.sqrt(x[0] ** 2 + x[1] ** 2 + 1e-12))
    beta1 = 2 * np.arcsin(np.sqrt(x[3]) ** 2 / np.sqrt(x[2] ** 2 + x[3] ** 2 + 1e-12))
    beta2 = 2 * np.arcsin(
        np.sqrt(x[2] ** 2 + x[3] ** 2) / np.sqrt(x[0] ** 2 + x[1] ** 2 + x[2] ** 2 + x[3] ** 2)
    )

    return np.array([beta2, -beta1 / 2, beta1 / 2, -beta0 / 2, beta0 / 2])

data = np.loadtxt("iris_classes1and2_scaled.txt")
X = data[:, 0:2]
print("First X sample (original)  :", X[0])

# pad the vectors to size 2^2 with constant values
padding = 0.3 * np.ones((len(X), 1))
X_pad = np.c_[np.c_[X, padding], np.zeros((len(X), 1))]
print("First X sample (padded)    :", X_pad[0])

# normalize each input
normalization = np.sqrt(np.sum(X_pad ** 2, -1))
X_norm = (X_pad.T / normalization).T
print("First X sample (normalized):", X_norm[0])

# angles for state preparation are new features
features = np.array([get_angles(x) for x in X_norm])
print("First features sample      :", features[0])

Y = (data[:, -1] + 1) / 2



def statepreparation(a, circuit, target):
    
    a = 2*a
    circuit.ry(a[0], target[0])
    
    circuit.cx(target[0], target[1])
    circuit.ry(a[1], target[1])
    circuit.cx(target[0], target[1])
    circuit.ry(a[2], target[1])

    circuit.x(target[0])
    circuit.cx(target[0], target[1])
    circuit.ry(a[3], target[1])
    circuit.cx(target[0], target[1])
    circuit.ry(a[4], target[1])
    circuit.x(target[0])
    
    return circuit

x = X_norm[0]
ang = get_angles(x)
q       = QuantumRegister(2)
c       = ClassicalRegister(1)
circuit = QuantumCircuit(q,c)
circuit = statepreparation(ang, circuit, [0,1])
circuit.draw(output='mpl')

def u_gate(param, circuit, target):
    '''Return the quantum circuit with u3 gate applied on qubit target with param as an iterable'''
    circuit.u3(param[0],param[1],param[2],target)
    return circuit

def cu_gate(param, circuit, control, target):
    '''Return the quantum circuit with cu3 gate applied on qubit target with param as an iterable wrt control'''
    circuit.cu3(param[0],param[1],param[2], control, target)
    return circuit

def circuit_block(param, circuit, target, same_order_x=True):
    '''Return the block applied on qubits target from the circuit circuit
    - param : array parameters for the two u gate
    - target : array of integer the numero of qubits for the u gates to be applied
    - if same_order_x == True : cx(target[0], target[1])
    else: cx(target[1], target[0])'''
    circuit = u_gate(param[0], circuit, target[0])
    circuit = u_gate(param[1], circuit, target[1])
    if same_order_x:
        circuit.cx(target[0], target[1])
    else:
        circuit.cx(target[1], target[0])
    return circuit

def c_circuit_block(param, circuit, control, target, same_order_x=True):
    '''Return the controlled block applied on qubits target from the circuit circuit
    - param : array parameters for the two u gate
    - target : array of integer the numero of qubits for the u gates to be applied
    - if same_order_x == True : cx(target[0], target[1])
    else: cx(target[1], target[0])'''
    circuit = cu_gate(param[0], circuit, control, target[0])
    circuit = cu_gate(param[1], circuit, control, target[1])
    if same_order_x:
        circuit.ccx(control, target[0], target[1])
    else:
        circuit.ccx(control, target[1], target[0])
    return circuit

def create_circuit(param, circuit, target):
    order = True
    for i in range(param.shape[0]):
        circuit = circuit_block(param[i], circuit, target, order)
        order = not order
    return circuit

def create_c_circuit(param, circuit, control, target):
    order = True
    for i in range(param.shape[0]):
        circuit = c_circuit_block(param[i], circuit, control, target, order)
        order = not order 
    return circuit

x = np.array([0.53896774, 0.79503606, 0.27826503, 0.0])
ang = get_angles(x)

params = np.array([[[np.pi/3,np.pi/3,np.pi/4],
                    [np.pi/6,np.pi/4,np.pi/6]],
                    [[np.pi/6,np.pi/4,np.pi/4],
                    [np.pi/3,np.pi/7,np.pi/6]],
                    [[np.pi/3,np.pi/3,np.pi/4],
                    [np.pi/6,np.pi/4,np.pi/6]],
                    [[np.pi/6,np.pi/4,np.pi/4],
                    [np.pi/3,np.pi/7,np.pi/6]]])

q       = QuantumRegister(2)
c       = ClassicalRegister(1)
circuit = QuantumCircuit(q,c)
circuit = statepreparation(ang, circuit, [0,1])

circuit = create_circuit(params, circuit, [0,1])

circuit.measure(0,c)

circuit.draw(output='mpl')

def execute_circuit(params, angles=None, x=None, use_angles=True, bias=0, shots=1000):
    if not use_angles:
        angles = get_angles(x)
    q       = QuantumRegister(2)
    c       = ClassicalRegister(1)
    circuit = QuantumCircuit(q,c)
    circuit = statepreparation(angles, circuit, [0,1])

    circuit = create_circuit(params, circuit, [0,1])
    circuit.measure(0,c)
    result = execute(circuit,backend,shots=shots).result()

    counts = result.get_counts(circuit)
    result=np.zeros(2)
    for key in counts:
        result[int(key,2)]=counts[key]
    result/=shots
    return result[1] + bias

execute_circuit(params, ang, bias=0.02)

def predict(probas):
    return (probas>=0.5)*1

def binary_crossentropy(labels, predictions):
    loss = 0
    for l, p in zip(labels, predictions):
        loss = loss - l * np.log(np.max([p,1e-8]))

    loss = loss / len(labels)
    return loss

def square_loss(labels, predictions):
    loss = 0
    for l, p in zip(labels, predictions):
        loss = loss + (l - p) ** 2

    loss = loss / len(labels)
    return loss 

def cost(params, features, labels):
    predictions = [execute_circuit(params, angles=f) for f in features]
    return binary_crossentropy(labels, predictions)

def accuracy(labels, predictions):

    loss = 0
    for l, p in zip(labels, predictions):
        if abs(l - p) < 1e-5:
            loss = loss + 1
    loss = loss / len(labels)

    return loss

def real(param1, param2, angles, shots=1000):
    """Returns Re{<circuit(param2)|sigma_z|circuit(param1)>}"""
    q = QuantumRegister(3)
    c = ClassicalRegister(1)
    circuit = QuantumCircuit(q,c)
    circuit.h(q[0])
    circuit = statepreparation(angles, circuit, [1,2])
    circuit = create_c_circuit(param1, circuit, 0, [1,2])
    circuit.cz(q[0], q[1])
    circuit.x(q[0])
    circuit = create_c_circuit(param2, circuit, 0, [1,2])
    circuit.x(q[0])
    circuit.h(q[0])
    circuit.measure(q[0],c)
    result = execute(circuit,backend,shots=shots).result()
    counts = result.get_counts(circuit)
    result=np.zeros(2)
    for key in counts:
        result[int(key,2)]=counts[key]
    result/=shots
    return (2*result[0]-1)

def imaginary(param1, param2, angles, shots=1000):
    """Returns Im{<circuit(param2)|sigma_z|circuit(param1)>}"""
    q = QuantumRegister(3)
    c = ClassicalRegister(1)
    circuit = QuantumCircuit(q,c)
    circuit.h(q[0])
    circuit = statepreparation(angles, circuit, [1,2])
    circuit = create_c_circuit(param1, circuit, 0, [1,2])
    circuit.cz(q[0], q[1])
    circuit.x(q[0])
    circuit = create_c_circuit(param2, circuit, 0, [1,2])
    circuit.x(q[0])
    circuit.u1(np.pi/2, q[0])
    circuit.h(q[0])
    circuit.measure(q[0],c)
    result = execute(circuit,backend,shots=shots).result()
    counts = result.get_counts(circuit)
    result=np.zeros(2)
    for key in counts:
        result[int(key,2)]=counts[key]
    result/=shots
    return -(2*result[0]-1)

def gradients(params, angles, label, bias=0):
    grads = np.zeros_like(params)
    imag = imaginary(params, params, angles)
    for i in range(params.shape[0]):
        for j in range(params.shape[1]):
            params_bis = np.copy(params)
            
            params_bis[i,j,0]+=np.pi
            grads[i,j,0] = -0.5 * real(params, params_bis, angles)
            params_bis[i,j,0]-=np.pi
            
            params_bis[i,j,1]+=np.pi
            grads[i,j,1] = 0.5 * (imaginary(params, params_bis, angles) - imag)
            params_bis[i,j,1]-=np.pi
            
            params_bis[i,j,2]+=np.pi
            grads[i,j,2] = 0.5 * (imaginary(params, params_bis, angles) - imag)
            params_bis[i,j,2]-=np.pi
    p = execute_circuit(params, angles, bias=bias) 
    grad_bias = (p - label) / (p * (1 - p))
    grads *= grad_bias
    return grads, grad_bias
            

np.random.seed(0)
num_data = len(Y)
num_train = int(0.75 * num_data)
index = np.random.permutation(range(num_data))
feats_train = features[index[:num_train]]
Y_train = Y[index[:num_train]]
feats_val = features[index[num_train:]]
Y_val = Y[index[num_train:]]

# We need these later for plotting
X_train = X[index[:num_train]]
X_val = X[index[num_train:]]


layers = 6
params_init = np.random.randn(layers,2,3) * 0.01
bias_init = 0.01
batch_size = 5
learning_rate = 0.01
momentum = 0.9

var = np.copy(params_init)
bias = bias_init
v = np.zeros_like(var)
v_bias = 0

for it in range(15):

    # Update the weights by one optimizer step
    batch_index = np.random.randint(0, num_train, (batch_size,))
    feats_train_batch = feats_train[batch_index]
    Y_train_batch = Y_train[batch_index]
    grads = np.zeros_like(var)
    grad_bias = 0
    var_corrected = var + momentum * v
    bias_corrected = bias + momentum * v_bias
    for j in range(batch_size):
        g, g_bias = gradients(var_corrected, feats_train_batch[j], Y_train_batch[j], bias)
        grads += g / batch_size
        grad_bias +=g_bias / batch_size
    
    v = momentum * v - learning_rate * grads
    v_bias = momentum * v_bias - learning_rate * grad_bias
    
    var += v
    bias += v_bias

    # Compute predictions on train and validation set
    probas_train = np.array([execute_circuit(var, angles=f, bias=bias) for f in feats_train])
    probas_val = np.array([execute_circuit(var, angles=f, bias=bias) for f in feats_val])
    predictions_train = predict(probas_train)
    predictions_val = predict(probas_val)

    # Compute accuracy on train and validation set
    acc_train = accuracy(Y_train, predictions_train)
    acc_val = accuracy(Y_val, predictions_val)

    print(
        "Iter: {:5d} | Loss: {:0.7f} | Acc train: {:0.7f} | Acc validation: {:0.7f} "
        "".format(it + 1, cost(var, features, Y), acc_train, acc_val)
    )

plt.figure()
cm = plt.cm.RdBu

# make data for decision regions
xx, yy = np.meshgrid(np.linspace(0.0, 1.5, 20), np.linspace(0.0, 1.5, 20))
X_grid = [np.array([x, y]) for x, y in zip(xx.flatten(), yy.flatten())]

# preprocess grid points like data inputs above
padding = 0.3 * np.ones((len(X_grid), 1))
X_grid = np.c_[np.c_[X_grid, padding], np.zeros((len(X_grid), 1))]  # pad each input
normalization = np.sqrt(np.sum(X_grid ** 2, -1))
X_grid = (X_grid.T / normalization).T  # normalize each input
features_grid = np.array(
    [get_angles(x) for x in X_grid]
)  # angles for state preparation are new features
predictions_grid = [execute_circuit(var, angles=f, bias=bias, shots=10000) for f in features_grid]
Z = np.reshape(predictions_grid, xx.shape)

# plot decision regions
cnt = plt.contourf(xx, yy, Z, levels=np.arange(0, 1.1, 0.1), cmap=cm, alpha=0.8, extend="both")
plt.contour(xx, yy, Z, levels=[0.0], colors=("black",), linestyles=("--",), linewidths=(0.8,))
plt.colorbar(cnt, ticks=[0, 0, 1])

# plot data
plt.scatter(
    X_train[:, 0][Y_train == 1],
    X_train[:, 1][Y_train == 1],
    c="b",
    marker="o",
    edgecolors="k",
    label="class 1 train",
)
plt.scatter(
    X_val[:, 0][Y_val == 1],
    X_val[:, 1][Y_val == 1],
    c="b",
    marker="^",
    edgecolors="k",
    label="class 1 validation",
)
plt.scatter(
    X_train[:, 0][Y_train == 0],
    X_train[:, 1][Y_train == 0],
    c="r",
    marker="o",
    edgecolors="k",
    label="class 0 train",
)
plt.scatter(
    X_val[:, 0][Y_val == 0],
    X_val[:, 1][Y_val == 0],
    c="r",
    marker="^",
    edgecolors="k",
    label="class 0 validation",
)

plt.legend()
plt.show()

