from sklearn import model_selection, datasets, svm
from qiskit import QuantumCircuit, Aer, IBMQ, QuantumRegister, ClassicalRegister
import numpy as np
import qiskit
import copy
import matplotlib.pyplot as plt
np.random.seed(42)

iris = datasets.load_iris()
X = iris.data[0:100]
Y = iris.target[0:100]
X_train, X_test, Y_train, Y_test = model_selection.train_test_split(X, Y, test_size=0.33, random_state=42)

N = 4

def feature_map(X, encoding_type):
    q = QuantumRegister(N)
    c = ClassicalRegister(1)
    qc = QuantumCircuit(q, c)
    if encoding_type == 'rx':
        for i, x in enumerate(X):
            qc.rx(x, i)
    elif encoding_type == 'ry':
        for i, x in enumerate(X):
            qc.ry(x, i)
    elif encoding_type == 'rz':
        for i, x in enumerate(X):
            qc.rz(x, i)    
    return qc, c

# feature map test
qc,c = feature_map(X[0], 'rx')
qc.draw('mpl')

def variational_circuit(qc, theta, rotation_type):
    
    for i in range(N-1):
        qc.cnot(i, i+1)
    qc.cnot(N-1, 0)
    
    if rotation_type == 'rx':
        for i in range(N):
            qc.rx(theta[i], i)
    elif rotation_type == 'ry':
        for i in range(N):
            qc.ry(theta[i], i)
    elif rotation_type == 'rz':
        for i in range(N):
            qc.rz(theta[i], i)
            
    return qc

qc = variational_circuit(qc, np.ones(4), 'ry')
qc.draw('mpl')

def quantum_nn(X, theta, shots, encoding_type, rotation_type, simulator=True):
    
    qc, c = feature_map(X, encoding_type)
    qc = variational_circuit(qc, theta, rotation_type)
    qc.measure(0, c)

    backend = Aer.get_backend('qasm_simulator')
    
    if simulator == False:
        provider = IBMQ.load_account()
        backend = provider.get_backend('ibm_oslo')
        
    job = qiskit.execute(qc, backend, shots=shots)
    result = job.result()
    counts = result.get_counts(qc)
    
    return counts['1']/shots


qc, c = feature_map(X[0], 'rz')
qc = variational_circuit(qc, np.ones(4), 'rz')
qc.measure(0, c)
backend = Aer.get_backend('qasm_simulator')
    
job = qiskit.execute(qc, backend, shots=5000)
result = job.result()
counts = result.get_counts(qc)
print(counts)

# quantum_nn test
quantum_nn(X[0], np.ones(4), 5000, 'rx', 'ry')

def loss(prediction, target):
    return(prediction-target)**2

# loss test
loss(quantum_nn(X[0], [0.9, 0.9, 0.9, 0.9], 5000, 'rx', 'ry'), Y[0])

def gradient(X, Y, theta, delta, shots, encoding_type, rotation_type):

    grad = []
    for i in range(len(theta)):
        
        dtheta = copy.copy(theta)
        dtheta[i] += delta
        
        prediction_1 = quantum_nn(X, dtheta, shots, encoding_type, rotation_type)
        prediction_2 = quantum_nn(X, theta, shots, encoding_type, rotation_type)
        
        grad.append((loss(prediction_1, Y) - loss(prediction_2, Y)) / delta)
        
    return np.array(grad)

gradient(X[0], Y[0], np.ones(4), 0.01, 5000, 'rx', 'ry')

def accuracy(X, Y, theta, shots, encoding_type, rotation_type):
    
    counter = 0
    for X_i, Y_i, in zip(X, Y):
        prediction = quantum_nn(X_i, theta, shots, encoding_type, rotation_type)
        
        if prediction < 0.5 and Y_i == 0:
            counter += 1
        elif prediction >= 0.5 and Y_i == 1:
            counter += 1
            
    return counter/len(Y)

# Accuracy test
accuracy(X_test, Y_test, [-1.5, 1.2, 0.6, -0.2], 1000, 'rx', 'ry')

def get_best_weight(accuracies, weights):
    return(weights[accuracies.index(max(accuracies))])


N = 4
def model(x_train, y_train, x_test, y_test, learning_rate, epochs, theta, shots, encoding_type, rotation_type, delta, optimizing_weights=False):
    loss_list = []
    acc_list = []
    weights_list = []

    print('Epoch \t Loss \t Training Accuracy \t theta')

    for i in range(epochs):
        loss_tmp = []


        for X_i, Y_i in zip(x_train, y_train):


            prediction = quantum_nn(X_i, theta, shots, encoding_type, rotation_type)
            loss_tmp.append(loss(prediction, Y_i))
            # update theta
            new_theta = theta - learning_rate * gradient(X_i, Y_i, theta, delta, shots, encoding_type, rotation_type)

            if loss(quantum_nn(X_i, new_theta, shots, encoding_type, rotation_type), Y_i) < loss(quantum_nn(X_i, theta, shots, encoding_type, rotation_type), Y_i):
                theta = new_theta

        loss_list.append(np.mean(loss_tmp))
        acc_list.append(accuracy(X_train, Y_train, theta, shots, encoding_type, rotation_type))
        weights_list.append(theta)
        
        print(f'{i} \t {round(loss_list[-1], 3)} \t {round(acc_list[-1], 3)} \t {theta} \t {len(weights_list)}')
    
    if optimizing_weights==True:
        acc = accuracy(x_test, y_test, get_best_weight(acc_list, weights_list), shots, encoding_type=encoding_gate, rotation_type=rotation_gate)
        return acc, get_best_weight(acc_list, weights_list), theta, loss_list, acc_list, weights_list
    else:
        acc = accuracy(x_test, y_test, theta, shots, encoding_type=encoding_gate, rotation_type=rotation_gate)
        return acc, theta, loss_list, acc_list, weights_list

gates = ['rx', 'ry', 'rz']
results = {}
for encoding_gate in gates:
    for rotation_gate in gates:
        circuit_type = encoding_gate + '-' + rotation_gate
        if circuit_type != 'rz-rz':
            print(f'Circuit type : {circuit_type} -> encoding : {encoding_gate} rotation : {rotation_gate}')
            acc, theta, loss_list, acc_list, weights_history = model(X_train, Y_train, X_test, Y_test, learning_rate=0.05, epochs=10, theta=np.ones(N), shots=5000, encoding_type=encoding_gate, rotation_type=rotation_gate, delta=0.01)
            results.update({circuit_type: {'accuracy': acc, 'theta': theta, 'loss_list':loss_list, 'acc_list':acc_list, 'weights_history_list':weights_history}})
            print(f" {circuit_type} accuracy : {results.get(circuit_type).get('accuracy')}")

gates = ['rx', 'ry', 'rz']
labels = []
accuracies = []
for encoding_gate in gates:
    for rotation_gate in gates:
        circuit_type = encoding_gate + '-' + rotation_gate
        if circuit_type != 'rz-rz':
            labels.append(circuit_type)
            accuracies.append(results.get(circuit_type).get('accuracy'))
            
plt.bar(labels, accuracies)
plt.show()

gates = ['rx', 'ry', 'rz']
labels = []
accuracies = []
for encoding_gate in gates:
    for rotation_gate in gates:
        circuit_type = encoding_gate + '-' + rotation_gate
        if circuit_type != 'rz-rz':
            plt.plot([x for x in range(1, 11)], results.get(circuit_type).get('acc_list'), label=circuit_type)
            
plt.xticks([x for x in range(1, 11)])
plt.yticks(np.arange (0, 1.1, 0.1))
plt.xlabel('Epochs')
plt.ylabel('Accuracy')
plt.title('Circuit Type accuracy effect')
plt.legend()
plt.show()


N = 4
def model_1(x_train, y_train, x_test, y_test, learning_rate, epochs, theta, shots, encoding_type, rotation_type, delta, optimizing_weights=False):
    loss_list = []
    acc_list = []
    weights_list = []
    loss_val = 1
    i = 0
    print('Epoch \t Loss \t Training Accuracy \t theta')

    while loss_val > 0.17:
        
        loss_tmp = []


        for X_i, Y_i in zip(x_train, y_train):


            prediction = quantum_nn(X_i, theta, shots, encoding_type, rotation_type)
            loss_tmp.append(loss(prediction, Y_i))
            # update theta
            new_theta = theta - learning_rate * gradient(X_i, Y_i, theta, delta, shots, encoding_type, rotation_type)

            if loss(quantum_nn(X_i, new_theta, shots, encoding_type, rotation_type), Y_i) < loss(quantum_nn(X_i, theta, shots, encoding_type, rotation_type), Y_i):
                theta = new_theta
        loss_val = np.mean(loss_tmp)
        loss_list.append(np.mean(loss_tmp))
        acc_list.append(accuracy(X_train, Y_train, theta, shots, encoding_type, rotation_type))
        weights_list.append(theta)
        
        print(f'{i} \t {round(loss_list[-1], 3)} \t {round(acc_list[-1], 3)} \t {theta} \t {len(weights_list)}')
        i = i + 1
    
    return theta, loss_list, acc_list, weights_list

theta_b, loss_list_b, acc_list_b, weights_history_b = model_1(X_train, Y_train, X_test, Y_test, learning_rate=0.05, epochs=20, theta=np.ones(N), shots=5000, encoding_type='ry', rotation_type='ry', delta=0.01)


N = 4
def model_2(x_train, y_train, x_test, y_test, learning_rate, epochs, theta, shots, encoding_type, rotation_type, delta, optimizing_weights=False):
    loss_list = []
    acc_list = []
    weights_list = []

    print('Epoch \t Loss \t Training Accuracy \t theta')

    for i in range(epochs):
        loss_tmp = []


        for X_i, Y_i in zip(x_train, y_train):


            prediction = quantum_nn(X_i, theta, shots, encoding_type, rotation_type)
            loss_tmp.append(loss(prediction, Y_i))
            # update theta
            new_theta = theta - learning_rate * gradient(X_i, Y_i, theta, delta, shots, encoding_type, rotation_type)

            if loss(quantum_nn(X_i, new_theta, shots, encoding_type, rotation_type), Y_i) < loss(quantum_nn(X_i, theta, shots, encoding_type, rotation_type), Y_i):
                theta = new_theta

        loss_list.append(np.mean(loss_tmp))
        acc_list.append(accuracy(X_train, Y_train, theta, shots, encoding_type, rotation_type))
        weights_list.append(theta)
        
        print(f'{i} \t {round(loss_list[-1], 3)} \t {round(acc_list[-1], 3)} \t {theta} \t {len(weights_list)}')
    
    return theta, loss_list, acc_list, weights_list

for learning_rate in [0.001, 0.01, 0.05, 0.1]:
    theta, loss_list, acc_list, weights_history = model_2(X_train, Y_train, X_test, Y_test, learning_rate=learning_rate, epochs=20, theta=np.ones(N), shots=5000, encoding_type='ry', rotation_type='ry', delta=0.01)
    acc = accuracy(X_test, Y_test, theta, 5000, encoding_type='ry', rotation_type='ry')
    results.update({learning_rate: {'accuracy': acc, 'theta': theta, 'loss_list':loss_list, 'acc_list':acc_list, 'weights_history_list':weights_history}})
    print(f" Learning rate : {learning_rate} accuracy : {results.get(learning_rate).get('accuracy')}")

for learning_rate in [0.001, 0.01, 0.05, 0.1]:
    plt.plot([x for x in range(1, 21)], results.get(learning_rate).get('acc_list'), label=learning_rate)
            
plt.xticks([x for x in range(1, 21)])
plt.yticks(np.arange (0, 1.1, 0.1))
plt.xlabel('eposhs')
plt.ylabel('Accuracy')
plt.title('learning rates - accuracy effect')
plt.legend()
plt.show()

n_shots_results = {}
for n_shots in [1000, 5000, 10000, 50000]:
    theta, loss_list, acc_list, weights_history = model_2(X_train, Y_train, X_test, Y_test, learning_rate=0.05, epochs=20, theta=np.ones(N), shots=n_shots, encoding_type='ry', rotation_type='ry', delta=0.01)
    acc = accuracy(X_test, Y_test, theta, n_shots, encoding_type='ry', rotation_type='ry')
    n_shots_results.update({n_shots: {'accuracy': acc, 'theta': theta, 'loss_list':loss_list, 'acc_list':acc_list, 'weights_history_list':weights_history}})
    print(f" Number of Shots : {n_shots} accuracy : {n_shots_results.get(n_shots).get('accuracy')}")

shots_list = [1000, 5000, 10000, 50000]
for n_shots in shots_list:
    plt.plot([x for x in range(1, 21)], n_shots_results.get(n_shots).get('acc_list'), label=n_shots)
            
plt.xticks([x for x in range(1, 21)])
plt.yticks(np.arange (0, 1.1, 0.1))
plt.xlabel('n_shots')
plt.ylabel('Accuracy')
plt.title('number of shots - accuracy effecta')
plt.legend()
plt.show()

delta_results = {}
for delta in [0.01, 0.05, 0.1, 1]:
    theta, loss_list, acc_list, weights_history = model_2(X_train, Y_train, X_test, Y_test, learning_rate=0.1, epochs=20, theta=np.ones(N), shots=1000, encoding_type='ry', rotation_type='ry', delta=delta)
    acc = accuracy(X_test, Y_test, theta, 1000, encoding_type='ry', rotation_type='ry')
    delta_results.update({delta: {'accuracy': acc, 'theta': theta, 'loss_list':loss_list, 'acc_list':acc_list, 'weights_history_list':weights_history}})
    print(f" Delta : {delta} accuracy : {delta_results.get(delta).get('accuracy')}")

for delta in [0.01, 0.05, 0.1, 1]:
    plt.plot([x for x in range(1, 21)], delta_results.get(delta).get('acc_list'), label=delta)
            
plt.xticks([x for x in range(1, 21)])
plt.yticks(np.arange (0, 1.1, 0.1))
plt.xlabel('eposhs')
plt.ylabel('Accuracy')
plt.title('Delta - accuracy effect')
plt.legend()
plt.show()

def feature_map_2(X, theta):
    n = 4

    qr = QuantumRegister(4)
    c = ClassicalRegister(1)
    qc = QuantumCircuit(qr,c )
    for i, x in enumerate(X):
        qc.ry(x, i)
    for i in range(n-1):
        qc.cx(i, i+1)

    qc.barrier()
    for i in range(4):
        qc.rz(theta[i], i)
    qc.barrier()

    for i in reversed(range(n-1)):
        qc.cx(i, i+1)
    qc.h(0)

    return qc, c

qc, c = feature_map_2(X[0], np.ones(4))
qc.draw('mpl')

def quantum_nn(X, theta, shots, encoding_type=None, rotation_type=None, simulator=True):
    
    qc, c = feature_map_2(X, theta)
    qc.measure(0, c)

    backend = Aer.get_backend('qasm_simulator')
    
    if simulator == False:
        provider = IBMQ.load_account()
        backend = provider.get_backend('ibm_oslo')
        
    job = qiskit.execute(qc, backend, shots=shots)
    result = job.result()
    counts = result.get_counts(qc)
    
    return counts['1']/shots

def accuracy(X, Y, theta, shots):
    
    counter = 0
    for X_i, Y_i, in zip(X, Y):
        prediction = quantum_nn(X_i, theta, shots)
        
        if prediction < 0.5 and Y_i == 0:
            counter += 1
        elif prediction >= 0.5 and Y_i == 1:
            counter += 1
            
    return counter/len(Y)

def gradient(X, Y, theta, delta, shots):

    grad = []
    for i in range(len(theta)):
        
        dtheta = copy.copy(theta)
        dtheta[i] += delta
        
        prediction_1 = quantum_nn(X, dtheta, shots)
        prediction_2 = quantum_nn(X, theta, shots)
        
        grad.append((loss(prediction_1, Y) - loss(prediction_2, Y)) / delta)
        
    return np.array(grad)

N = 4
def model_5(x_train, y_train, x_test, y_test, learning_rate, epochs, theta, shots, delta, optimizing_weights=False):
    loss_list = []
    acc_list = []
    weights_list = []

    print('Epoch \t Loss \t Training Accuracy \t theta')

    for i in range(epochs):
        loss_tmp = []


        for X_i, Y_i in zip(x_train, y_train):


            prediction = quantum_nn(X_i, theta, shots)
            loss_tmp.append(loss(prediction, Y_i))
            # update theta
            new_theta = theta - learning_rate * gradient(X_i, Y_i, theta, delta, shots)

            if loss(quantum_nn(X_i, new_theta, shots), Y_i) < loss(quantum_nn(X_i, theta, shots), Y_i):
                theta = new_theta

        loss_list.append(np.mean(loss_tmp))
        acc_list.append(accuracy(X_train, Y_train, theta, shots))
        weights_list.append(theta)
        
        print(f'{i} \t {round(loss_list[-1], 3)} \t {round(acc_list[-1], 3)} \t {theta} \t {len(weights_list)}')
    
    return theta, loss_list, acc_list, weights_list

theta, loss_list, acc_list, weights_list = model_5(X_train, Y_train, X_test, Y_test, 0.01, 30, [1, 1, 1, 1], 5000, 0.01)

accuracy(X_test, Y_test, theta, 5000)

plt.plot([x for x in range(1, 31)], acc_list)
            
plt.xticks([x for x in range(1, 31)])
plt.yticks(np.arange (0, 1.1, 0.1))
plt.xlabel('eposhs')
plt.ylabel('Accuracy')
plt.title('Accuracy')
plt.legend()
plt.show()

from qiskit.circuit.library import ZZFeatureMap

num_features = X_train.shape[1]

feature_map = ZZFeatureMap(feature_dimension=num_features, reps=1)
feature_map.decompose().draw(output="mpl", fold=20)

from qiskit.circuit.library import RealAmplitudes

ansatz = RealAmplitudes(num_qubits=num_features, reps=3)
ansatz.decompose().draw(output="mpl", fold=20)

from qiskit.algorithms.optimizers import COBYLA

optimizer = COBYLA(maxiter=100)

from qiskit.primitives import Sampler

sampler = Sampler()

from matplotlib import pyplot as plt
from IPython.display import clear_output

objective_func_vals = []
plt.rcParams["figure.figsize"] = (12, 6)


def callback_graph(weights, obj_func_eval):
    clear_output(wait=True)
    objective_func_vals.append(obj_func_eval)
    plt.title("Objective function value against iteration")
    plt.xlabel("Iteration")
    plt.ylabel("Objective function value")
    plt.plot(range(len(objective_func_vals)), objective_func_vals)
    plt.show()

import time
from qiskit_machine_learning.algorithms.classifiers import VQC

vqc = VQC(
    sampler=sampler,
    feature_map=feature_map,
    ansatz=ansatz,
    optimizer=optimizer,
    callback=callback_graph,
)

# clear objective value history
objective_func_vals = []

start = time.time()
vqc.fit(X_train, Y_train)
elapsed = time.time() - start

print(f"Training time: {round(elapsed)} seconds")

train_score_q4 = vqc.score(X_train, Y_train)
test_score_q4 = vqc.score(X_test, Y_test)


print(f"Quantum QNN on the training dataset: {accuracy(X_train, Y_train, theta, 5000):.2f}")
print(f"Quantum QNN on the test dataset:     {accuracy(X_test, Y_test, theta, 5000):.2f}")
print(f"Quantum VQC on the training dataset: {train_score_q4:.2f}")
print(f"Quantum VQC on the test dataset:     {test_score_q4:.2f}")

