from svm.datasets import *
from qiskit.aqua.input import ClassificationInput
from qiskit.aqua import run_algorithm
from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit
from qiskit import Aer, execute
from qiskit.aqua.components.optimizers.cobyla import COBYLA
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

# size of training data set
training_size = 20

# size of test data set
test_size = 10

# dimension of data sets
n = 2

# construct training and test data
# set the following flag to True for the first data set and to False for the second dataset
use_adhoc_dataset = True
if use_adhoc_dataset:
    # first (artifical) data set to test the classifier
    training_input, test_input, class_labels = \
            ad_hoc_data(training_size=training_size, test_size=test_size, n=n, gap=0.3, plot_data=True)
else:
    # second data set to test the classifier
    training_input, test_input, class_labels = \
            Wine(training_size=training_size, test_size=test_size, n=n, plot_data=True)


def feature_map(x, q):
    
    # initialize quantum circuit
    qc = QuantumCircuit(q)
    
    # apply hadamards and U_Phi twice
    for _ in range(2):
        
        # apply the hadamard and Z-rotatiion to all qubits
        for i in range(x.shape[0]):
            qc.h(q[i])
            qc.rz(2 * x[i], q[i])
        
        # apply the two qubit gate
        qc.cx(q[0], q[1])
        qc.rz(2*(np.pi-x[0])*(np.pi-x[1]), q[1])
        qc.cx(q[0], q[1])

    # return quantum circuit
    return qc

# initialize quantum register
num_qubits = 2
qr = QuantumRegister(num_qubits)

# initialize test data (x1, x2)
data = np.asarray([1.5, 0.3])

# get quantum circuit
qc_feature_map = feature_map(data, qr)
# simulate using local statevector simulator
backend = Aer.get_backend('statevector_simulator')

job_sim = execute(qc_feature_map, backend)
sim_results = job_sim.result()
print('simulation: ', sim_results)
print('statevector: ', np.round(sim_results.get_statevector(), decimals=4))

# draw circuit
qc_feature_map.draw(output='mpl', plot_barriers=False)

def variational_form(q, params, depth):
    
    # initialize quantum circuit
    qc = QuantumCircuit(q)
    
    # first set of rotations
    param_idx = 0
    for qubit in range(2):
        qc.ry(params[param_idx], q[qubit])
        qc.rz(params[param_idx+1], q[qubit])
        param_idx += 2

    # entangler blocks and succeeding rotations
    for block in range(depth):
        qc.cz(q[0], q[1])
        for qubit in range(2):
            qc.ry(params[param_idx], q[qubit])
            qc.rz(params[param_idx+1], q[qubit])
            param_idx += 2
    
    # return quantum circuit
    return qc

# initialize quantum register
num_qubits = 2
qr = QuantumRegister(num_qubits)

# set depth, i.e. number of entangler blocks and rotations (after the initial rotation)
depth = 2
params = [0.5,-0.3, 0.2, 0.6]*(depth+1)
qc_variational_form = variational_form(qr, params, depth)

# simulate using local statevector simulator
job_sim = execute(qc_variational_form, backend)
sim_results = job_sim.result()
print('simulation: ', sim_results)
print('statevector: ', np.round(sim_results.get_statevector(), decimals=4))

# draw circuit
qc_variational_form.draw(output='mpl', plot_barriers=False)

def assign_label(bit_string, class_labels):
    hamming_weight = sum([int(k) for k in list(bit_string)])
    is_odd_parity = hamming_weight & 1
    if is_odd_parity:
        return class_labels[1]
    else:
        return class_labels[0]

# assigns a label 
assign_label('01', class_labels)

def return_probabilities(counts, class_labels):
    shots = sum(counts.values())
    result = {class_labels[0]: 0, 
              class_labels[1]: 0}
    for key, item in counts.items():
        label = assign_label(key, class_labels)
        result[label] += counts[key]/shots
    return result

return_probabilities({'00' : 10, '01': 10}, class_labels)

def classifier_circuit(x, params, depth):
    
    q = QuantumRegister(2)
    c = ClassicalRegister(2)
    qc = QuantumCircuit(q, c)

    qc_feature_map = feature_map(x, q)
    qc_variational_form = variational_form(q, params, depth)

    qc += qc_feature_map
    qc += qc_variational_form

    qc.measure(q, c)
    
    return qc

def classify(x_list, params, class_labels, depth=2, shots=100):
    
    qc_list = []
    for x in x_list:
        qc = classifier_circuit(x, params, depth)
        qc_list += [qc]
    qasm_backend = Aer.get_backend('qasm_simulator')    
    jobs = execute(qc_list, qasm_backend, shots=shots)
    
    probs = []
    for qc in qc_list:
        counts = jobs.result().get_counts(qc)
        prob = return_probabilities(counts, class_labels)
        probs += [prob]
            
    return probs

# draw classifier circuit
qc = classifier_circuit(np.asarray([0.5, 0.5]), params, depth)


# classify test data point (using random parameters constructed earlier)
x = np.asarray([[0.5, 0.5]])
classify(x, params, class_labels, depth)

qc.draw(output='mpl', plot_barriers=False)

def cost_estimate_sigmoid(probs, expected_label):

    p = probs.get(expected_label)
    sig = None
    if np.isclose(p, 0.0):
        sig = 1
    elif np.isclose(p, 1.0):
        sig = 0
    else:
        denominator = np.sqrt(2*p*(1-p))
        x = np.sqrt(200)*(0.5-p)/denominator
        sig = 1/(1+np.exp(-x))

    return sig

x = np.linspace(0, 1, 20)
y = [cost_estimate_sigmoid({'A': x_, 'B': 1-x_}, 'A') for x_ in x]
plt.plot(x, y)
plt.xlabel('Probability of assigning the correct class')
plt.ylabel('Cost value')
plt.show()

def cost_function(training_input, class_labels, params, depth=2, shots=100, print_value=False):
    
    # map training input to list of labels and list of samples
    cost = 0
    training_labels = []
    training_samples = []
    for label, samples in training_input.items():
        for sample in samples:
            training_labels += [label]
            training_samples += [sample]
        
    # classify all samples
    probs = classify(training_samples, params, class_labels, depth)
    
    # evaluate costs for all classified samples
    for i, prob in enumerate(probs):
        cost += cost_estimate_sigmoid(prob, training_labels[i])
    cost /= len(training_samples)
    
    # print resulting objective function
    if print_value:
        print('%.4f' % cost)
    
    # return objective value
    return cost

cost_function(training_input, class_labels, params, depth)

# set depth of variational form
depth = 2

# set number of shots to evaluate the classification circuit
shots = 100

# setup the optimizer
optimizer = COBYLA()

# define objective function for training
objective_function = lambda params: cost_function(training_input, class_labels, params, depth, shots, print_value=True)

# randomly initialize the parameters
init_params = 2*np.pi*np.random.rand(num_qubits*(depth+1)*2)

# train classifier
opt_params, value, _ = optimizer.optimize(len(init_params), objective_function, initial_point=init_params)

# print results
print()
print('opt_params:', opt_params)
print('opt_value: ', value)

# collect coordinates of test data
test_label_0_x = [x[0] for x in test_input[class_labels[0]]]
test_label_0_y = [x[1] for x in test_input[class_labels[0]]]
test_label_1_x = [x[0] for x in test_input[class_labels[1]]]
test_label_1_y = [x[1] for x in test_input[class_labels[1]]]

# initialize lists for misclassified datapoints
test_label_misclassified_x = []
test_label_misclassified_y = []

# evaluate test data
for label, samples in test_input.items():
    
    # classify samples 
    results = classify(samples, opt_params, class_labels, depth, shots=shots)
    
    # analyze results
    for i, result in enumerate(results):
        
        # assign label
        assigned_label = class_labels[np.argmax([p for p in result.values()])]
        print('----------------------------------------------------')
        print('Data point:    ', samples[i])
        print('Label:         ', label)
        print('Assigned:      ', assigned_label)
        print('Probabilities: ', result)
        
        if label != assigned_label:
            print('Classification:', 'INCORRECT')
            test_label_misclassified_x += [samples[i][0]]
            test_label_misclassified_y += [samples[i][1]]
        else:
            print('Classification:', 'CORRECT')
        
# compute fraction of misclassified samples
total = len(test_label_0_x) + len(test_label_1_x)
num_misclassified = len(test_label_misclassified_x)
print()
print(100*(1-num_misclassified/total), "% of the test data was correctly classified!")

# plot results
plt.figure()
plt.scatter(test_label_0_x, test_label_0_y, c='b', label=class_labels[0], linewidths=5)
plt.scatter(test_label_1_x, test_label_1_y, c='g', label=class_labels[1], linewidths=5)
plt.scatter(test_label_misclassified_x, test_label_misclassified_y, linewidths=20, s=1, facecolors='none', edgecolors='r')
plt.legend()
plt.show()

