from sklearn import datasets
import numpy as np
import matplotlib.pyplot as plt

from qiskit.quantum_info import DensityMatrix

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

from qCircuit import createInputCircuit, createModelCircuit, createNoisyModelCircuit
from qLipschitz import QLipschitz
from qDistance import d, D

import warnings
warnings.filterwarnings('ignore')

X_train, Y_train = datasets.make_circles(n_samples=500, noise=0.09, factor=0.4)
X_test, Y_test = datasets.make_circles(n_samples=200, 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)

X, Y = datasets.load_iris(return_X_y=True)

X_train_0 = X[0:30]
X_train_1 = X[50:80]
X_train_iris = np.vstack((X_train_0, X_train_1))

Y_train_0 = Y[0:30]
Y_train_1 = Y[50:80]
Y_train_iris = np.vstack((Y_train_0, Y_train_1)).flatten()

X, Y = datasets.load_iris(return_X_y=True)

X_test_0 = X[30:50]
X_test_1 = X[80:100]
X_test_iris = np.vstack((X_test_0, X_test_1))

Y_test_0 = Y[30:50]
Y_test_1 = Y[80:100]
Y_test_iris = np.vstack((Y_test_0, Y_test_1)).flatten()

X_train_iris = (X_train_iris - X_train_iris.min(axis=0)) / (X_train_iris.max(axis=0) - X_train_iris.min(axis=0))
X_test_iris = (X_test_iris - X_test_iris.min(axis=0)) / (X_test_iris.max(axis=0) - X_test_iris.min(axis=0))

import qiskit
from qiskit import Aer, execute

qasm_backend = Aer.get_backend('qasm_simulator')

params = [
    [[-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]]
]

params_iris = [
    [[-3.42142600e-02, -5.45047534e-03,  8.11905347e-01],
     [-1.52484152e+00,  1.49884676e+00, -1.25408626e-02],
     [-1.89515860e-03,  1.35760410e-02,  8.22999582e-03],
     [ 1.39102114e+00,  4.19729865e-01,  1.60000380e-04]],

    [[ 7.88431068e-01, -8.86177264e-01,  1.33830291e-02],
     [-3.71228143e-03, -1.12994101e-02, -1.27897783e-02],
     [ 9.45954683e-03, -3.34659883e-03,  1.17217829e-02],
     [ 1.98256181e-02, -1.07358054e-02,  7.53621360e-03]]
]

# Test 1: createModelCircuit
createModelCircuit(params).draw(output='mpl')

def predict(model, data):
    qubits_num = len(data)
    input_circuit = createInputCircuit(data)
    qcircuit = input_circuit.compose(model)
    
    # the measurement
    qcircuit.measure(list(range(qubits_num)), list(range(qubits_num)))

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

    p1 = 0
    for i in range(2 ** (qubits_num - 1)):
        comb = str(bin(i + 2 ** (qubits_num - 1))[2:])
        p1 += counts.get(comb, 0)
        
    p1 /= shots
    
    if p1 > 0.5:
        return 1
    else:
        return 0

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

def evaluate(model, X_test, Y_test):
    acc = accuracy(model, X_test, Y_test)

    print("FINAL ACCURACY: {:0.2f}%".format(acc * 100))

# Test 2: evaluate noiseless model

test2_circuit = createModelCircuit(params_iris)

evaluate(test2_circuit, X_train_iris, Y_train_iris)
evaluate(test2_circuit, X_test_iris, Y_test_iris)

# Test 3: create noisy model circuit - circle

test3_p = 0.001
test3_circuit = createNoisyModelCircuit(params, test3_p, 'phase_flip')
test3_circuit.draw(output='mpl')

# Test 4: evaluate noisy model - circle

test4_p = 0.001
models_type = [
    'noiseless', 'bit_flip', 'phase_flip', 'depolarizing'
]
test4_circuits = [
    createModelCircuit(params),
    createNoisyModelCircuit(params, test4_p, 'bit_flip'),
    createNoisyModelCircuit(params, test4_p, 'phase_flip'),
    createNoisyModelCircuit(params, test4_p, 'depolarizing')
]

for i, model_type in enumerate(models_type):
    print(f'Evaluate <{model_type}>: ')
    evaluate(test4_circuits[i], X_test2, Y_test)
    print()

# Test 5: calculate lipschitz constant

print('Classifier for centric circle: ')
test5_1_circuit = createModelCircuit(params)
test5_1_qlipschitz = QLipschitz(test5_1_circuit, [0, 1])
print(test5_1_qlipschitz.constant)

print('\nClassifier for Iris: ')
test5_2_circuit = createModelCircuit(params_iris)
test5_2_qlipschitz = QLipschitz(test5_2_circuit, [0, 1])
print(test5_2_qlipschitz.constant)

# Test 6-1: calculate lipschitz constant in noisy model

test6_p = [0.0001, 0.001, 0.05, 0.1]
models_type = [
    'noiseless', 'bit_flip', 'phase_flip', 'depolarizing'
]

print('Classifier for centric circle: ')
for p in test6_p:
    test6_1_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):
        test6_1_qlipschitz = QLipschitz(test6_1_circuits[i], [0, 1])
        print(f'{model_type} model: {test6_1_qlipschitz.constant[0]}')
    print()

# Test 6-2: calculate lipschitz constant in noisy model

test6_p = [0.0001, 0.001, 0.05, 0.1]
models_type = [
    'noiseless', 'bit_flip', 'phase_flip', 'depolarizing'
]

print('Classifier for Iris: ')
for p in test6_p:
    test6_2_circuits = [
        createModelCircuit(params_iris),
        createNoisyModelCircuit(params_iris, p, 'bit_flip'),
        createNoisyModelCircuit(params_iris, p, 'phase_flip'),
        createNoisyModelCircuit(params_iris, p, 'depolarizing')
    ]
    
    print(f'p: {p}')
    for i, model_type in enumerate(models_type):
        test6_2_qlipschitz = QLipschitz(test6_2_circuits[i], [0, 1])
        print(f'{model_type} model: {test6_2_qlipschitz.constant[0]}')
    print()

# Test 7-1: find proper epsilon - circle

import random

# D(rou, sigma) <= epsilon
indices = []
while len(indices) < 50:
    a = random.randint(0, len(X_train2) - 1)
    b = random.randint(0, len(X_train2) - 1)
    if a != b:
        indices.append([a, b])

distances = []
for pair in indices:
    rou = getDensityMatrix(createInputCircuit(X_train2[pair[0]]))
    sigma = getDensityMatrix(createInputCircuit(X_train2[pair[1]]))
    
    distances.append(D(rou, sigma))
    
print('find proper epsilon - centric circle')
print(f'Mean distance: {sum(distances) / len(distances)}')
print(f'Max distance: {max(distances)}')
print(f'Min distance: {min(distances)}')

# Test 7-2: find proper epsilon - iris

# D(rou, sigma) <= epsilon
indices = []
while len(indices) < 50:
    a = random.randint(0, len(X_train_iris) - 1)
    b = random.randint(0, len(X_train_iris) - 1)
    if a != b:
        indices.append([a, b])

distances = []
for pair in indices:
    rou = getDensityMatrix(createInputCircuit(X_train_iris[pair[0]]))
    sigma = getDensityMatrix(createInputCircuit(X_train_iris[pair[1]]))
    
    distances.append(D(rou, sigma))

print('find proper epsilon - centric iris')
print(f'Mean distance: {sum(distances) / len(distances)}')
print(f'Max distance: {max(distances)}')
print(f'Min distance: {min(distances)}')

# Test 8-1: find proper delta - circle

test8_1_circuit = createModelCircuit(params)
test8_1_qlipschitz = QLipschitz(test8_1_circuit, [0, 1])
test8_1_A = test8_1_qlipschitz.A

# d(A(rou), A(sigma)) <= delta
indices = []
while len(indices) < 50:
    a = random.randint(0, len(X_train2) - 1)
    b = random.randint(0, len(X_train2) - 1)
    if a != b:
        indices.append([a, b])

distances = []
for pair in indices:
    rou = getDensityMatrix(createInputCircuit(X_train2[pair[0]]))
    sigma = getDensityMatrix(createInputCircuit(X_train2[pair[1]]))
    
    distances.append(d(test8_1_A, rou, sigma))

print('find proper delta - circle')
print(f'Mean distance: {sum(distances) / len(distances)}')
print(f'Max distance: {max(distances)}')
print(f'Min distance: {min(distances)}')

# Test 8-2: find proper delta - circle

test8_2_circuit = createModelCircuit(params_iris)
test8_2_qlipschitz = QLipschitz(test8_2_circuit, [0, 1])
test8_2_A = test8_2_qlipschitz.A

# d(A(rou), A(sigma)) <= delta
indices = []
while len(indices) < 50:
    a = random.randint(0, len(X_train_iris) - 1)
    b = random.randint(0, len(X_train_iris) - 1)
    if a != b:
        indices.append([a, b])

distances = []
for pair in indices:
    rou = getDensityMatrix(createInputCircuit(X_train_iris[pair[0]]))
    sigma = getDensityMatrix(createInputCircuit(X_train_iris[pair[1]]))
    
    distances.append(d(test8_2_A, rou, sigma))

print('find proper delta - iris')
print(f'Mean distance: {sum(distances) / len(distances)}')
print(f'Max distance: {max(distances)}')
print(f'Min distance: {min(distances)}')

# Test 9-1: fairness verifying - circle

test9_1_circuit = createNoisyModelCircuit(params, p=0.001, errorType='bit_flip')
test9_1_qlipschitz = QLipschitz(test9_1_circuit, [0, 1])
print(f'Lipschitz constant: {test9_1_qlipschitz.constant[0]}\n')

test9_1_fairness_params = [
    # epsilon, delta
    [1, 1], [0.9, 0.89], [0.6, 0.58], [0.6, 0.59]
]

for epsilon, delta in test9_1_fairness_params:
    flag, kernel = test9_1_qlipschitz.fairVeriQ(epsilon, delta)
    
    print(f'A is ({epsilon},{delta})-fair: <{flag}>')
    if not flag:
        print(f' - bias kernel pair is: {kernel}')
    print()

# Test 9-2: fairness verifying - iris

test9_2_circuit = createNoisyModelCircuit(params_iris, p=0.001, errorType='phase_flip')
test9_2_qlipschitz = QLipschitz(test9_2_circuit, [0, 1])
print(f'Lipschitz constant: {test9_2_qlipschitz.constant[0]}\n')

test9_2_fairness_params = [
    # epsilon, delta
    [1, 0.99], [0.8, 0.79], [0.5, 0.3]
]

for epsilon, delta in test9_2_fairness_params:
    flag, kernel = test9_1_qlipschitz.fairVeriQ(epsilon, delta)
    
    print(f'A is ({epsilon},{delta})-fair: <{flag}>')
    if not flag:
        print(f' - bias kernel pair is: {kernel}')
    print()

# Test 10: generate bias pair

test10_circuit = createNoisyModelCircuit(params_iris, p=0.01, errorType='depolarizing')
test10_qlipschitz = QLipschitz(test10_circuit, [0, 1])

print(f'Lipschitz constant: {test10_qlipschitz.constant[0]}\n')

flag, kernel = test10_qlipschitz.fairVeriQ(epsilon=0.9, delta=0.3)
print(f'fairness: {flag}')
assert (not flag)

test10_sigma = getDensityMatrix(createInputCircuit(X_train_iris[11]))
test10_bias_pair = test10_qlipschitz.generateBiasPair(test10_sigma, kernel, epsilon=0.9)
test10_rou_psi, test10_rou_phi = test10_bias_pair
# print(test10_bias_pair)

print(f'epsilon=0.9, delta=0.3')
print('D: ', D(test10_rou_psi, test10_rou_phi))
print('d: ', d(test10_qlipschitz.A, test10_rou_psi, test10_rou_phi))

