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.09, factor=0.4)
X_test, Y_test = datasets.make_circles(n_samples=50, noise=0.09, factor=0.4)

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)

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

        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))

@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
import numpy as np
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)

vqc = VQC()

# vqc.params = np.array([[[-1.18991246e-01, -8.69694713e-01,  1.43722811e-03],
#       [ 3.47641545e+00, -4.73632073e-02,  1.97135609e+00],
#       [ 1.95775707e+00, -1.12139160e-02, -6.21796144e-03]],

#      [[ 2.02460380e-02, -2.25797547e+00,  5.14234265e-03],
#       [-1.71552299e-02,  2.46283604e-03, -1.03805722e-02],
#       [-2.37244982e-03, -3.35799404e-03,  6.10191152e-03]]])

vqc.fit(X_train2, Y_train, epoch=10)

print(cost(vqc.bestparams, X_train2, Y_train))
print(cost(vqc.bestparams, X_test2, Y_test))

vqc.score(X_test2, Y_test)

def createCircuit(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])
        
    return qcircuit

qcircuit = createCircuit(vqc.bestparams, X_train2[0])
qcircuit.draw(output='mpl')

def measure(qcircuit, label):
    # 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)

    print(f'simulation results: \n{counts}')
    p1 = (counts.get('100', 0) + counts.get('111', 0) + counts.get('101', 0) + counts.get('110',0)) / shots
    print(f'p1: {p1}')
    print(f'Expected label: {label}')

example = [3, 30, 35, 70]

for index in example:
    qcircuit = createCircuit(vqc.bestparams, X_train2[index])
    measure(qcircuit, Y_train[index])
    print()

# origin data
plt.scatter(X_train2[:, 0], X_train2[:, 1], c=Y_train)
plt.show()

predictions = [ predict(vqc.bestparams, data) for data in X_test2 ]
Y_ = [ Y_test[i] if Y_test[i] == predictions[i] else 2  for i in range(len(Y_test)) ]

plt.scatter(X_test2[:, 0], X_test2[:, 1], c=Y_)
plt.show()

train_params(X_train2[10:50], Y_train[10:50])

