from sklearn import datasets
import pennylane as qml
import numpy as np
from pennylane.optimize import AdamOptimizer
import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings('ignore')

X_train, Y_train = datasets.make_circles(n_samples=200, noise=0.06, factor=0.6)
X_test, Y_test = datasets.make_circles(n_samples=80, noise=0.06, factor=0.6)

print('Training Data:')
print(X_train[:10])
print('\nTraining Label:')
print(Y_train[:10])

plt.scatter(X_train[:, 0], X_train[:, 1], c=Y_train)
plt.show()

print(f'Total datas: {len(X_train)}')
print(f'labels: {list(set(Y_train))}')
print(f'label 0 datas: {len(list(filter(lambda x: x==0, Y_train)))}')
print(f'label 1 datas: {len(list(filter(lambda x: x==1, Y_train)))}')

Z_train = (X_train[:, 0] ** 2 + X_train[:, 1] ** 2) ** 0.5
temp = np.zeros((X_train.shape[0], len(X_train[0]) + 1))
temp[:, :-1] = X_train
temp[:, -1] = Z_train
X_train2 = temp

Z_test = (X_test[:, 0] ** 2 + X_test[:, 1] ** 2) ** 0.5
temp = np.zeros((X_test.shape[0], len(X_test[0]) + 1))
temp[:, :-1] = X_test
temp[:, -1] = Z_test
X_test2 = temp

print('Solved Training Data:')
print(X_train2[:10])

def plot3D(X, Y):
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(
        X[:, 0], X[:, 1], X[:, 2], 
        zdir='z', s=30, c=Y, depthshade=True
    )
    plt.show()

plot3D(X_train2, Y_train)

print(f'Total datas: {len(X_train2)}')
print(f'labels: {list(set(Y_train))}')
print(f'label 0 datas: {len(list(filter(lambda x: x==0, Y_train)))}')
print(f'label 1 datas: {len(list(filter(lambda x: x==1, Y_train)))}')

qubits_num = 3
layers_num = 2

dev = qml.device("default.qubit", wires=qubits_num)

@qml.qnode(dev)
def circuit(params, data):
    angles = [ i  * np.pi for i in data ]

    for i in range(qubits_num):
        qml.RX(angles[i], wires=i)
        qml.Rot( *params[0, i], wires=i )

    qml.CZ(wires=[1, 0])
    qml.CZ(wires=[1, 2])
    qml.CZ(wires=[0, 2])

    for i in range(qubits_num):
        qml.Rot( *params[1, i], wires=i )
    
    # PauliZ measure => 1 -> |0>  while -1 -> |1>
    return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)), qml.expval(qml.PauliZ(2))

import qiskit
from qiskit import QuantumCircuit
from qiskit import Aer, execute

unitary_backend = Aer.get_backend('unitary_simulator')
qasm_backend = Aer.get_backend('qasm_simulator')

def predict(params, data):
    qcircuit = QuantumCircuit(3, 3)
    qubits = qcircuit.qubits

    for i, d in enumerate(data):
        qcircuit.rx(d * np.pi, qubits[i])
        
    for i in range(qubits_num):
        qcircuit.u3(*params[0][i], qubits[i])

    qcircuit.cz(qubits[0], qubits[1])
    qcircuit.cz(qubits[1], qubits[2])
    qcircuit.cz(qubits[0], qubits[2])

    for i in range(qubits_num):
        qcircuit.u3(*params[1][i], qubits[i])
    
    # the measurement
    qcircuit.measure([0, 1, 2], [0, 1, 2])

    # job execution
    shots = 1000
    job_sim = execute(qcircuit, qasm_backend, shots=shots)
    result_sim = job_sim.result()
    counts = result_sim.get_counts(qcircuit)

    p1 = (counts.get('100', 0) + counts.get('111', 0) + counts.get('101', 0) + counts.get('110',0)) / shots
    
    if p1 > 0.5:
        return 1
    else:
        return 0

def cost(weights, datas, labels):
    loss = 0
    
    for i, data in enumerate(datas):
        # like [-1, 1, 1]
        measured = circuit(weights, data)

        p = measured[0]
        if labels[i] == 0:
            loss += (1 - p) ** 2
        else:
            loss += (-1 - p) ** 2
        
    return loss / len(datas)

def accuracy(weights, datas, labels):
    predictions = [ predict(weights, data) for data in datas ]
    
    acc = 0
    for i, p in enumerate(predictions):
        if p == labels[i]:
            acc += 1
    
    return acc / len(predictions)

class VQC:
    def __init__(self):
        # 3 => U3(theta, phi, lambda)
        self.params = (0.01 * np.random.randn(layers_num, qubits_num, 3))
        
        self.bestparams = self.params
        self.bestcost = 10
        
        self.weights = []
        self.costs = []
        self.accuracies = []
        
    def fit(self, X_train, Y_train, *, epoch=300, stepsize=0.1):
        self.opt = AdamOptimizer(stepsize)
        batch_size = 30

        for turn in range(epoch):
            # Update the weights by one optimizer step
            batch_index = np.random.randint(0, len(X_train), (batch_size,))
            X_train_batch = X_train[batch_index]
            Y_train_batch = Y_train[batch_index]
            self.params = self.opt.step(lambda v: cost(v, X_train_batch, Y_train_batch), self.params)

            cost_now = cost(self.params, X_train, Y_train)
            acc_now = accuracy(self.params, X_train, Y_train) 

            if cost_now < self.bestcost:
                self.bestcost = cost_now
                self.bestparams = self.params
            
            self.weights.append(self.params)
            self.costs.append(cost_now)
            self.accuracies.append(acc_now)
            
            print(
                "Turn: {:5d} | Cost: {:0.7f} | Accuracy: {:0.2f}% ".format(
                turn + 1, cost_now, acc_now * 100
            ))
            
    def score(self, X_test, Y_test):
        predictions = [ predict(self.bestparams, data) for data in X_test ]
        acc = accuracy(self.bestparams, X_test, Y_test)
        
        print("FINAL ACCURACY: {:0.2f}%".format(acc * 100))

def train_params(datas, labels, *, epoch=30):
    alpha_grid = [0.001, 0.05, 0.1, 0.125, 0.25, 0.35, 0.5, 0.75, 1, 1.25, 1.5, 1.75, 2, 2.25, 2.5, 2.75, 3]
#     alpha_grid = [0.05, 0.15, 0.25, 0.5]
#     alpha_grid = [0.05, 0.075, 0.1, 0.125, 0.25, 0.5, 0.75, 1]
    alpha_best = None
    loss_best = 1000
    batch_size = 30
    
    for alpha in alpha_grid:
        print(f'alpha: {alpha}: ')
        vqc = VQC()
        vqc.fit(datas, labels, epoch=epoch, stepsize=alpha)
        loss = vqc.bestcost
        print(f' - finished. (loss {loss})\n')

        if loss < loss_best:
            loss_best = loss
            alpha_best = alpha

    return {
        'alpha': alpha_best,
        'loss': loss_best 
    }

train_params()

vqc = VQC()

vqc.fit(X_train2, Y_train, epoch=300, stepsize=)

