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=400, noise=0.06, factor=0.7)
X_test, Y_test = datasets.make_circles(n_samples=100, noise=0.06, factor=0.7)

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

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

    for i in range(qubits_num):
        qml.RX(angles[i], wires=i)
        qml.Rot( *(params_fixed[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_training[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 cost2(params_fixed, params_training, datas, labels):
    loss = 0
    
    for i, data in enumerate(datas):
        # like [-1, 1, 1]
        measured = circuit2(params_fixed, params_training, 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.random(size=(layers_num, qubits_num, 3))
        self.params = (0.01 * np.random.randn(layers_num, qubits_num, 3))
        self.bestparams = self.params
        self.bestcost = 10
        self.bestaccuracy = 0
        
        self.weights = []
        self.costs = []
        self.accuracies = []
        
    def fit(self, X_train, Y_train, *, epoch=300, stepsize=0.125, batch_size=20, fixed=False):
        self.opt = AdamOptimizer(stepsize)

        for turn in range(epoch):
            # Update the weights by one optimizer step
            if batch_size != -1:
                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]
                if fixed:
                    params_fixed = self.params[0]
                    params_training = self.params[1]
                    self.params[1] = self.opt.step(lambda v: cost2(params_fixed, v, X_train_batch, Y_train_batch), params_training)
                else:
                    self.params = self.opt.step(lambda v: cost(v, X_train_batch, Y_train_batch), self.params)
            else:
                if fixed:
                    params_fixed = self.params[0]
                    params_training = self.params[1]
                    self.params[1] = self.opt.step(lambda v: cost2(params_fixed, v, X_train, Y_train), params_training)
                else:
                    self.params = self.opt.step(lambda v: cost(v, X_train, Y_train), 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))

from qiskit.quantum_info import DensityMatrix

def getDensityMatrix(circuit):
    return DensityMatrix(circuit).data

from scipy.linalg import sqrtm

def D(rou, sigma):
    A = rou - sigma
    A_ = sqrtm( np.dot( A.conj().T, A ) )

    return 0.5 * np.linalg.norm( np.trace(A_) )

def params_loss_distance(params, data1, data2, label1, label2):
    u3Circuits = []
    for i in range(qubits_num):
        circuit = QuantumCircuit(1, 1)
        circuit.u3(*( params[i] ), 0)
        u3Circuits.append(circuit)
        
    circuit1 = QuantumCircuit(qubits_num, qubits_num)
    for qubit in range(qubits_num):
        circuit1.rx(data1[qubit] * np.pi, qubit)
        
    circuit2 = QuantumCircuit(qubits_num, qubits_num)
    for qubit in range(qubits_num):
        circuit2.rx(data2[qubit] * np.pi, qubit)
    
    matrix1 = getDensityMatrix(circuit1)
    matrix2 = getDensityMatrix(circuit2)
    d = D(matrix1, matrix2)
    
    if label1 == label2:
        loss = d ** 2
    else:
        loss = (1 - d) ** 2
    
    return loss

def params_loss_gradient(params, datas, labels, *, loss_function, epsilon=0.001):
    length = len(datas)
    derivates = np.zeros(len(params))
    loss_origin = 0
    
    # calculate origin loss
    for i in range(length // 2):
        j = i + length // 2
        loss_origin += loss_function(params, datas[i], datas[j], labels[i], labels[j])
    loss_origin /= length // 2
    
    for i in range(len(params)):
        params[i] += epsilon
        loss = 0
        
        for j in range(length // 2):
            k = j + length // 2
            loss += loss_function(params, datas[j], datas[k], labels[j], labels[k])
        loss /= length // 2
        
        derivates[i] = (loss - loss_origin) / epsilon
        params[i] -= epsilon
    
    return loss_origin, derivates

import random

def sub_train_params(datas, labels, *, epoch, alpha, batch_size, loss_function):
#     params = 0.01 * np.random.random(size=(1, qubits_num, 3))[0]
    params = (0.01 * np.random.randn(1, qubits_num, 3))[0]
#     params = params.reshape(qubits_num, 3)
    print(f'Training at alpha: {alpha} ')
    
    momentum, s = 0, 0
    belta1 = 0.9; belta2 = 0.999
    e = 1e-3 # 1e-8
    lastloss = 100
    
    for turn in range(epoch):
        if turn % 10 == 9: 
            print(f'  -- training: ({turn + 1}/{epoch})')
        
        batch_index = random.sample(range(0, len(datas)), batch_size)
#         batch_index = range(0, 300)
        loss, grad = params_loss_gradient(params, datas[batch_index], labels[batch_index],
                                          loss_function=loss_function)

        if abs(loss - lastloss) < 1e-4:
            break
        lastloss = loss

        # Adam Optimizer
        momentum = belta1 * momentum + (1 - belta1) * grad
        s = belta2 * s + (1 - belta2) * (grad ** 2)
        m_ = momentum / (1 - belta1 ** (turn + 1))
        s_ = s / (1 - belta2 ** (turn + 1))

        # update parameters
        params = params - alpha * (m_ / ( s_ ** 0.5 + e ))
    
    return loss, params

def train_params(datas, labels, *, epoch=30, loss_function=params_loss_distance):
#     alpha_grid = [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; params_best = None; loss_best = 1000
    batch_size = 30
    
    for alpha in alpha_grid:
        loss, params = sub_train_params(datas, labels, 
            alpha=alpha, batch_size=batch_size, epoch=epoch, loss_function=loss_function)
        print(f' - finished. (loss {loss})\n')

        if loss < loss_best:
            loss_best = loss
            alpha_best = alpha
            params_best = params

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

result = train_params(X_train2, Y_train, epoch=300)

print(f'best alpha: {result["alpha"]}')
print(f'best loss: {result["loss"]}')
print(f'best params: {result["params"]}')

vqc = VQC()
stepsize = 0.001

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

params_origin = vqc.bestparams
params_origin

print(cost(params_origin, X_train2, Y_train))
print(cost(params_origin, X_test2, Y_test))

vqc.score(X_test2, Y_test)

vqc = VQC()
vqc.params[0] = result['params']
stepsize = result['alpha']
# stepsize = 0.125

vqc.fit(X_train2, Y_train, epoch=10, stepsize=stepsize, fixed=False)

params_init = vqc.bestparams
params_init

print(cost(params_init, X_train2, Y_train))
print(cost(params_init, X_test2, Y_test))

vqc.score(X_test2, Y_test)

params_test = vqc.bestparams
params_test

vqc = VQC()
print(f'training with alpha: {result["alpha"]}')
vqc.params[0] = result['params']
# vqc.params[1] = result['params']
stepsize = result['alpha'] / 100
# stepsize = 

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

print(cost(params_test, X_train2, Y_train))
print(cost(params_test, X_test2, Y_test))

vqc.score(X_test2, Y_test)

params_u3 = vqc.bestparams
print(params_u3)

def createInputCircuit(data):
    qubits_num = len(data)
    qcircuit = QuantumCircuit(qubits_num, qubits_num)
    qubits = qcircuit.qubits

    for i, d in enumerate(data):
        qcircuit.rx(d * np.pi, qubits[i])
        
    return qcircuit

def createModelCircuit(params):
    qubits_num = len(params[0])
    qcircuit = QuantumCircuit(qubits_num, qubits_num)
    qubits = qcircuit.qubits

    for i in range(qubits_num):
        qcircuit.u3(*params[0][i], qubits[i])

    for i in range(qubits_num - 1):
        qcircuit.cz(qubits[i], qubits[i + 1])
    qcircuit.cz(qubits[0], qubits[qubits_num - 1])

    for i in range(qubits_num):
        qcircuit.u3(*params[1][i], qubits[i])
        
    return qcircuit

def createCircuit(params, data):
    input_circuit = createInputCircuit(data)
    model_circuit = createModelCircuit(params)
    full_circuit = input_circuit.compose(model_circuit)
    
    return full_circuit

from qiskit.providers.aer.noise import NoiseModel
from qiskit.providers.aer.utils import insert_noise
from qiskit.providers.aer.noise import pauli_error, depolarizing_error, mixed_unitary_error
from qiskit.circuit.library import XGate, IGate, ZGate
from qiskit.quantum_info import DensityMatrix

def createNoiseModel(p, errorType):
    # QuantumError objects
    if errorType == 'bit_flip':
        error = pauli_error([('X', p), ('I', 1 - p)])
    elif errorType == 'phase_flip':
        error = pauli_error([('Z', p), ('I', 1 - p)])
    elif errorType == 'depolarizing':
        error = depolarizing_error(p, num_qubits=1)
    elif errorType == 'mixed':
        error = mixed_unitary_error([(XGate(), p/2), (ZGate(), p/2), (IGate(), 1 - p)])
        
    ## two-qubits quantumError objects 
    if errorType == 'depolarizing':
        error_2qubits = depolarizing_error(p, num_qubits=2)
    else:
        error_2qubits = error.tensor(error)
        
    # Add errors to noise model
    noise_model = NoiseModel()
    
    noise_model.add_all_qubit_quantum_error(error, ['u3'])
    noise_model.add_all_qubit_quantum_error(error_2qubits, ['cz'])
    
    return noise_model

def createNoisyModelCircuit(params, p, errorType, *, model_circuit=None):
    noise_model = createNoiseModel(p, errorType)
    if model_circuit == None:
        model_circuit = createModelCircuit(params)
    
    return insert_noise(model_circuit, noise_model)

from functools import reduce

Dag = lambda matrix: matrix.conj().T
Kron = lambda *matrices: reduce(np.kron, matrices)

psi_0 = np.array([1.0, 0.0])
psi_1 = np.array([0.0, 1.0])
I = np.eye(2)

M_0 = psi_0.reshape([2, 1]) @ psi_0.reshape([1, 2]).conj()
M_1 = psi_1.reshape([2, 1]) @ psi_1.reshape([1, 2]).conj()

def getMeasurements(qubits_num):
    measurement_0 = [M_0]
    measurement_1 = [M_1]
    
    for i in range(qubits_num - 1):
        measurement_0.append(I)
        measurement_1.append(I)
        
    return [
        Kron(*measurement_0),
        Kron(*measurement_1)
    ]

def powerSets(items):
    N = len(items)
    combs = []
    
    for i in range(2 ** N):
        comb = []
        for j in range(N):
            if (i >> j) % 2:
                comb.append(items[j])
        combs.append(comb)
    
    return combs 

class Algorithm:
    def __init__(self, model_circuit, measurements, outputs):
        # DensityMatrix of model
        self.E = getDensityMatrix(model_circuit)
        
        # Measurements
        self.M = dict()
        for index, output in enumerate(outputs):
            self.M[output] = measurements[index]
        
        # Outputs
        self.O = outputs
        self.O_ = powerSets(outputs)

def d(A, rou, sigma):
    distance = 0
    
    for output in A.O:
        trace = np.trace(
            Dag(A.M[output]) @ A.M[output] @ 
            (A.E @ (rou - sigma) @ Dag(A.E))
        )
        
        distance += np.linalg.norm(trace)
    
    return distance / 2

def Lipschitz(A):
    E, M, O, O_ = A.E, A.M, A.O, A.O_
    
    # Step 1: Calculate W_i
    W = dict()
    for i in O:
        W[i] = Dag(E) @ Dag(M[i]) @ M[i] @ E
    
    # Step 2: Calculate K_star
    K_star = 0; vectors = [None, None]
    M_star = np.zeros(E.shape)
    
    for S in O_:
        if len(S) == 0:
            continue
            
        # calculate M_S = Σ Wi
        M_S = np.zeros(E.shape).astype('complex64')
        for i in S:
            M_S += W[i]
        
        # calculate eigenvalues and eigenvectors of M_S
        eigenvalues, eigenvectors = np.linalg.eigh(M_S)
        min_index = np.where(eigenvalues == eigenvalues.min())
        max_index = np.where(eigenvalues == eigenvalues.max())
        
        # calculate K_S
        K_S = np.linalg.norm(eigenvalues[max_index][0] - eigenvalues[min_index][0])
        
        if K_S > K_star:
            K_star = K_S
            vectors[0] = eigenvectors.T[max_index][0]
            vectors[1] = eigenvectors.T[min_index][0]
            
    return K_star, np.array(vectors)

def FairVeriQ(A, epsilon, delta):
    # epsilon <= 1 and delta > 0
    K_star, kernel = Lipschitz(A)
    
    if delta >= K_star * epsilon:
        return True, None
    else:
        return False, kernel

def generateBiasPair(sigma, kernel, epsilon):
    psi, phi = kernel
    size = len(psi)
    psi = psi.reshape(size, 1) @ Dag(psi.reshape(size, 1))
    phi = phi.reshape(size, 1) @ Dag(phi.reshape(size, 1))
    
    rou_psi = epsilon * psi + (1 - epsilon) * sigma
    rou_phi = epsilon * phi + (1 - epsilon) * sigma
    
    return np.array([
        rou_psi, rou_phi
    ])

X_train2[0]

# params trained by noiesless environment
vqc.bestparams

# calculate lipschitz constant in noisy model
probabilities = [0.0001, 0.001, 0.05, 0.1]
models_type = [
    'noiseless', 'bit_flip', 'phase_flip', 'depolarizing'
]

print('Classifier for centric circle: ')
params = vqc.bestparams
for p in probabilities:
    circuits = [
        createModelCircuit(params),
        createNoisyModelCircuit(params, p, 'bit_flip'),
        createNoisyModelCircuit(params, p, 'phase_flip'),
        createNoisyModelCircuit(params, p, 'depolarizing')
    ]
    
    print(f'p: {p}')
    for i, model_type in enumerate(models_type):
        noisy_A = Algorithm(circuits[i], getMeasurements(qubits_num=3), [0, 1])
        print(f'{model_type} model: {Lipschitz(noisy_A)[0]}')
    print()

def createParamTrainingCircuit(params):
    qubits_num = len(params) // 3
    qcircuit = QuantumCircuit(qubits_num, qubits_num)
    qubits = qcircuit.qubits

    for i in range(qubits_num):
        qcircuit.u3(*params[i*3: (i+1)*3], qubits[i])

    return qcircuit

def params_loss_noisemodel_distance_factory(p, noiseType):
    
    def params_loss_noisemodel_distance(params, data1, data2, label1, label2):
        u3Circuit = QuantumCircuit(qubits_num, qubits_num)
        qubits = u3Circuit.qubits
        for i in range(qubits_num):
            u3Circuit.u3(*params[i*3: (i+1)*3], qubits[i])
        
        inputCircuit1 = QuantumCircuit(qubits_num, qubits_num)
        qubits = inputCircuit1.qubits
        for i, d in enumerate(data1):
            inputCircuit1.rx(d * np.pi, qubits[i])

        inputCircuit2 = QuantumCircuit(qubits_num, qubits_num)
        qubits = inputCircuit2.qubits
        for i, d in enumerate(data1):
            inputCircuit2.rx(d * np.pi, qubits[i])


        # calculate distance
        matrix_1 = getDensityMatrix(circuit1)
        matrix_2 = getDensityMatrix(circuit2)

        d = D(matrix_1, matrix_2)
        if label1 == label2:
            return d ** 2
        else:
            return (1 - d) ** 2
    
    return params_loss_noisemodel_distance

params_loss_noisemodel_distance = params_loss_noisemodel_distance_factory(0.001, 'bit_flip')
result_noise = train_params(X_train2, Y_train, epoch=200, loss_function=params_loss_noisemodel_distance)

print(f'best alpha: {result_noise["alpha"]}')
print(f'best loss: {result_noise["loss"]}')
print(f'best params: {result_noise["params"]}')

vqc = VQC()
vqc.params[0] = result_noise['params']
stepsize = result_noise['alpha'] / 10

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

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

vqc.score(X_test2, Y_test)

# calculate lipschitz constant in noisy model
probabilities = [0.0001, 0.001, 0.05, 0.1]
models_type = [
    'noiseless', 'bit_flip', 'phase_flip', 'depolarizing'
]

print('Classifier for centric circle: ')
params = vqc.bestparams
for p in probabilities:
    circuits = [
        createModelCircuit(params),
        createNoisyModelCircuit(params, p, 'bit_flip'),
        createNoisyModelCircuit(params, p, 'phase_flip'),
        createNoisyModelCircuit(params, p, 'depolarizing')
    ]
    
    print(f'p: {p}')
    for i, model_type in enumerate(models_type):
        noisy_A = Algorithm(circuits[i], getMeasurements(qubits_num=3), [0, 1])
        print(f'{model_type} model: {Lipschitz(noisy_A)[0]}')
    print()

circuit_ = createNoisyModelCircuit(params, 0.01, 'mixed')
circuit_.draw(output='mpl')

