import numpy as np
import random
import matplotlib.pyplot as plt
from sklearn.neighbors import KNeighborsClassifier
from sklearn.neighbors import NearestCentroid
from IPython.display import clear_output
import pandas as pd
# =======================================================
import math
import time
from qiskit import Aer, QuantumCircuit, execute
from qiskit.extensions import UnitaryGate
from qiskit.providers.jobstatus import JobStatus


# Use Aer's qasm_simulator by default
backend = Aer.get_backend('qasm_simulator')

# generate toy data: random N x N board
def generate_board(N):
    '''
    Generates a random board of dimensions N x N
    Inputs:
    N          := 1D dimension of board, s.t. board.shape = (N,N)

    Outputs:
    board      := flattened N x N array
    winner     := red wins == 1, blue wins == -1
    sum_red    := sum of red pieces
    sum_blue   := sum of blue pieces
    '''
    board=np.zeros((N*N))
    pieces=random.sample(range(1,N*N),1)[0]
    inds = random.sample(range(0,N*N),pieces) # pick random location to place pieces
    ratio = 1
    while ratio == 1: # safeguard to remove possiblility of ties
        for n in range(pieces):
            board[inds[n]]=random.sample([-1,1],1)[0] # make space blue or red
        sum_red = np.sum(board==1) # sum of 1 values (i.e., red pieces)
        sum_blue = np.sum(board==-1) # sum of -1 values (i.e., blue pieces)
        ratio = (sum_red+1)/(sum_blue+1) # find ratio
        if ratio > 1: # red wins
            winner = 1
        elif ratio < 1: # blue wins
            winner = -1
        else:
            pass
    return board, winner, sum_red, sum_blue


def generate_data(D,N,train_split):
    '''
    Generates labelled data arrays + images.

    Inputs:
    D                    := desired number of data points
    N                    := 1D dimension of board, s.t. board.shape = (N,N)
    train_split          := split of training/testing data, e.g., 0.9 => 90% training, 10% testing

    Outputs:
    training_data        := split training data, data array: [# red pieces, # blue pieces, winner] dims: (D*train_split) x 3
    training_images      := split training data, corresponding board images; flattened dims: (D*train_split) x N**2
    testing_data         := split testing data, data array: [# red pieces, # blue pieces, winner] dims: (D*(1-train_split)) x 3
    testing_images       := split testing data, corresponding board images; flattened dims: (D*(1-train_split)) x N**2
    '''
    data = np.zeros((D,3)) # create array to populate each row with x1: black sum, x2: white sum, label: ratio
    images = np.zeros((D,N**2)) # create array to populate D rows with flattened N x N images in the columns
    for n in range(D):
        board,winner,sum_red,sum_blue = generate_board(N=N)
        data[n,[0,1,2]]=[sum_red,sum_blue,winner]
        images[n,:]=board
    training_data=data[:int(D*train_split),:] # labelled training data
    training_images=images[:int(D*train_split),:] # training images
    testing_data=data[int(D*train_split):,:] # labelled testing data
    testing_images=images[int(D*train_split):,:] # testing images
    return training_data, training_images, testing_data, testing_images

# gameboard parameters
N=8 # 1D of gameboard. Game board => N x N dimensions
# generate training/testing data for the classical and quantum classifiers
training_data, training_images, testing_data, testing_images =generate_data(D=300,N=8,train_split=0.9)


def nCen_train_test(training_data, testing_data, N, shrink, plot_testing=False):
    '''
    Train Nearest Centroid Neighbors algorithm. Plot results and accuracy. Output the best nCen.

    Inputs:
    training_data    := split training data, data array: [# red pieces, # blue pieces, winner] dims: (D*train_split) x 3
    testing_data     := split testing data, data array: [# red pieces, # blue pieces, winner] dims: (D*(1-train_split)) x 3
    N                := 1D dimension of board, s.t. board.shape = (N,N); used here for normalizing the data
    shrink           := list of shrink values to test, e.g., shrink=[0.0,0.1,0.5,0.9]
    plot_testing     := True or False value; plots the testing dataset with classifier bounds

    Outputs:
    Plots the nearest centroid decision boundary, training data, test data, and accuracy plots
    nCen             := nearest centroid object with the highest accuracy
    '''

    # Assign training/test data
    training_data[:,[0,1]]=training_data[:,[0,1]]/(N**2) # standardize data. Max pieces is N*N
    testing_data[:,[0,1]]=testing_data[:,[0,1]]/(N**2) # standardize data. Max pieces is N*N
    X_train = training_data[:,[0,1]] # training features
    y_train = training_data[:,2] # training labels
    X_test = testing_data[:,[0,1]] # testing features
    y_test = testing_data[:,2] # testing labels
    # =============================================================================

    # run accuracy with training data for k=k neighbors
    acc_list=[]
    nCen_list=[]
    for i in range(len(shrink)):

        # run kNN classification
        nCen_list.append(NearestCentroid(shrink_threshold=shrink[i])) # create k-NN object
        nCen_list[i].fit(X_train, y_train) # fit classifier decision boundary to training data
        # =============================================================================

        # plot classification boundary with training data
        meshres=50
        RED,BLUE=np.meshgrid(np.linspace(np.min([np.min(X_train[:,0]),np.min(X_test[:,0])]),np.max([np.max(X_train[:,0]),np.max(X_test[:,0])]),meshres),
                            np.linspace(np.min([np.min(X_train[:,1]),np.min(X_test[:,1])]),np.max([np.max(X_train[:,1]),np.max(X_test[:,1])]),meshres))
        boundary=nCen_list[i].predict(np.c_[RED.ravel(),BLUE.ravel()])
        plt.figure()
        plt.contourf(RED,BLUE,boundary.reshape(RED.shape),cmap='bwr',alpha=0.3)
        plt.scatter(X_train[:,0],X_train[:,1],c=y_train,cmap='bwr',edgecolor='k')
        plt.xlabel('# red pieces\n[normalized]')
        plt.ylabel('# blue pieces\n[normalized]')
        cbar=plt.colorbar()
        cbar.ax.set_ylabel('Winner')
        plt.title('Centroid Classification, shrink = '+str(shrink[i]))
        plt.show()
    
        # plot testing data
        if plot_testing == True:
            plt.figure()
            plt.contourf(RED,BLUE,boundary.reshape(RED.shape),cmap='bwr',alpha=0.3)
            testt=plt.scatter(X_test[:,0],X_test[:,1],c=y_test,cmap='bwr',edgecolor='k')
            plt.xlabel('# red pieces\n[normalized]')
            plt.ylabel('# blue pieces\n[normalized]')
            plt.title('Testing Data, n = '+str(X_test.shape[0]))
            cbar=plt.colorbar(testt)
            cbar.ax.set_ylabel('Winner')
            plt.show()
        else:
            pass
        # =============================================================================

        # calculate accuracy with training data set
        predicted=nCen_list[i].predict(X_test) # calculate predicted label from training data
        acc=np.sum(predicted==y_test)/len(y_test) # calculate accuracy of where prediction is correct vs incorrect relative to actual labels
        acc_list.append(acc)
        print('shrink = ',shrink[i],'\nTraining Accuracy =',acc*100,'%\n','='*40,'\n\n')
    plt.figure()
    plt.plot(shrink,acc_list,marker='o')
    plt.xlabel('Centroid shrinking threshold')
    plt.xticks(shrink)
    plt.ylabel('Accuracy')
    plt.show()
    best_nCen=nCen_list[np.argmax(acc_list)]
    print('best shrink = ',shrink[np.argmax(acc_list)])
    return best_nCen

# generate best nearest centroid neighbors
best_nCen = nCen_train_test(training_data=training_data, testing_data=testing_data, N=N, shrink=list(np.arange(0.0,0.9,0.1)),plot_testing=True)

def player_prediction_classical(best_model,player_data):
    '''
    Outputs the classically predicted result of the 5 player games, not within the training dataset.

    Inputs:
    best_model  := NearestCentroid object
    player_data := generated data array: [# red pieces, # blue pieces, winner] dims: 5 x 3

    Outputs:
    cCompData   := list of 0 or 1 values. 1 == red wins, 0 == blue wins
    '''
    X_player_data=player_data[:,[0,1]]
    cCompData = best_model.predict(X_player_data)
    cCompData[cCompData==-1.]=0
    return list(cCompData.astype(int))

# Unit-normalized vector helper
def normalize(x):
    return x / np.linalg.norm(x)

# Since in a space R^n there are n-1 RBS gates,
# in a 2-dimensional dataset (R^2), there is 1
# RBS parameter angle, as defined:
def theta_in_2d(x):
    epsilon = math.pow(10, -10)
    return math.acos((x[0] + epsilon) / (np.linalg.norm(x) + epsilon)) + epsilon


def estimation_circuit(x_theta, y_theta):
    # As defined in Johri et al. (III.A.1), an iRBS is functionally
    # identical in the distance estimation circuit to the RBS gate;
    # thus, the iRBS gate is defined in the QuantumCircuit class to
    # be used in the estimation circuit.

    # Unitary implementation (throws IonQ backend error)
    def irbs_unitary(theta):
        return UnitaryGate([[1, 0, 0, 0],
                            [0, np.cos(theta), complex(0,-np.sin(theta)), 0],
                            [0, complex(0, -np.sin(theta)), np.cos(theta), 0],
                            [0, 0, 0, 1]], label=f'iRBS({np.around(theta, 3)})')

    # Circuit implementation from figure 7 in Johri et al.
    def irbs(qc, theta, qubits):
        qc.rz(np.pi/2, qubits[1])
        qc.cx(qubits[1], qubits[0])
        qc.ry(np.pi/2 - 2 * theta, qubits[1])
        qc.rz(-np.pi/2, qubits[0])
        qc.cx(qubits[0], qubits[1])
        qc.ry(2 * theta - np.pi/2, qubits[1])
        qc.cx(qubits[1], qubits[0])
        qc.rz(-np.pi / 2, qubits[0])

    qc = QuantumCircuit(2,2)
    qc.x(0)
    irbs(qc, x_theta, [0, 1])
    irbs(qc, y_theta, [0, 1])
    qc.measure([0, 1], [0, 1])
    return qc

# Returns estimated distance as defined in Johri et al. (Equation 3)
def probability_uncorrected(result):
    counts = result.get_counts()
    shots = np.sum(list(counts.values()))
    if '10' in counts:
        return counts["10"]/shots
    else:
        return 0

def probability_corrected(result):
    counts = result.get_counts()
    corrected_counts = {}

    # Discard incorrect values (00, 11)
    if '10' in counts:
        corrected_counts['10'] = counts['10']
    if '01' in counts:
        corrected_counts['01'] = counts['01']

    shots = np.sum(list(corrected_counts.values()))
    if '10' in counts:
        return corrected_counts["10"]/shots
    else:
        return 0

def euclidian_distance(x, y, corrected=False):
    x_norm = np.linalg.norm(x)
    y_norm = np.linalg.norm(y)
    x_square_norm = x_norm**2
    y_square_norm = y_norm**2

    x_theta = theta_in_2d(x)
    y_theta = theta_in_2d(y)


    ec = estimation_circuit(x_theta, y_theta)

    if backend.name() == "qasm_simulator":
        job = execute(ec, backend, shots=1000)
    else: # IonQ
        job = backend.run(ec, shots=1000)

    job_id = job.job_id()

    while job.status() is not JobStatus.DONE:
        print(job.status())
        time.sleep(1)
  
    result = job.result()

    if corrected:
        probability = probability_corrected(result)
    else:
        probability = probability_uncorrected(result)

    normalized_inner_product = math.sqrt(probability)

    return math.sqrt(x_square_norm + y_square_norm - 2*x_norm*y_norm*normalized_inner_product)


import random

# generate toy data: random N x N board
def generate_board(N):
    '''
    Generates a random board of dimensions N x N
    Inputs:
    N          := 1D dimension of board, s.t. board.shape = (N,N)

    Outputs:
    board      := flattened N x N array
    winner     := red wins == 1, blue wins == -1
    sum_red    := sum of red pieces
    sum_blue   := sum of blue pieces
    '''
    board=np.zeros((N*N))
    pieces=random.sample(range(1,N*N),1)[0]
    inds = random.sample(range(0,N*N),pieces) # pick random location to place pieces
    ratio = 1
    while ratio == 1: # safeguard to remove possiblility of ties
        for n in range(pieces):
            board[inds[n]]=random.sample([-1,1],1)[0] # make space blue or red
        sum_red = np.sum(board==1) # sum of 1 values (i.e., red pieces)
        sum_blue = np.sum(board==-1) # sum of -1 values (i.e., blue pieces)
        ratio = sum_red/sum_blue # find ratio
        if ratio > 1: # red wins
            winner = 1
        elif ratio < 1: # blue wins
            winner = -1
        else:
            pass
    return board, winner, sum_red, sum_blue

class QuantumNearestCentroid:
    def __init__(self):
        self.centroids = {}
  
    def centroid(self, x):
        return np.mean(x, axis=0)

    def fit(self, x, y):
        organized_values = {}
        for i, v in enumerate(y):
            if v in organized_values:
                organized_values[v].append(x[i])
        else:
            organized_values[v] = [x[i]]
        for v in organized_values:
            self.centroids[v] = self.centroid(organized_values[v])
        return self.centroids

    def predict(self, x):
        min_dist = np.inf
        closest_centroid = None
        for key, centroid in self.centroids.items():
            res = euclidian_distance(x, centroid)
            if res < min_dist:
                min_dist = res
                closest_centroid = key
        return closest_centroid

def player_prediction_quantum(best_model,player_data):
    '''
    Outputs the quantum predicted result of the 5 player games, not within the training dataset.

    Inputs:
    best_model  := NearestCentroid object
    player_data := generated data array: [# red pieces, # blue pieces, winner] dims: 5 x 3

    Outputs:
    cCompData   := list of 0 or 1 values. 1 == red wins, 0 == blue wins
    '''
    X_player_data=player_data[:,[0,1]]
    qCompData = []
    for i in range(5):
        qCompData.append(best_model.predict(X_player_data[i]))
    qCompData[qCompData==-1.]=0
    return list(qCompData)

def quantum_train_test(training_data, testing_data, N, plot_testing=False):
    '''
    Train Nearest Centroid Neighbors algorithm. Plot results and accuracy. Output the best nCen.

    Inputs:
    training_data    := split training data, data array: [# red pieces, # blue pieces, winner] dims: (D*train_split) x 3
    testing_data     := split testing data, data array: [# red pieces, # blue pieces, winner] dims: (D*(1-train_split)) x 3
    N                := 1D dimension of board, s.t. board.shape = (N,N); used here for normalizing the data
    shrink           := list of shrink values to test, e.g., shrink=[0.0,0.1,0.5,0.9]
    plot_testing     := True or False value; plots the testing dataset with classifier bounds

    Outputs:
    Plots the nearest centroid decision boundary, training data, test data, and accuracy plots
    nCen             := nearest centroid object with the highest accuracy
    '''

    # Assign training/test data
    X_train = training_data[:,[0,1]] # training features
    y_train = training_data[:,2] # training labels
    X_test = testing_data[:,[0,1]] # testing features
    y_test = testing_data[:,2] # testing labels
    # =============================================================================

    model = QuantumNearestCentroid()

    centroid = model.fit(X_train.tolist(), y_train.flatten().tolist())
    print(centroid)

    meshres=30

    RED,BLUE=np.meshgrid(np.linspace(np.min([np.min(X_train[:,0]),np.min(X_test[:,0])]),np.max([np.max(X_train[:,0]),np.max(X_test[:,0])]),meshres),
                      np.linspace(np.min([np.min(X_train[:,1]),np.min(X_test[:,1])]),np.max([np.max(X_train[:,1]),np.max(X_test[:,1])]),meshres))
    print(len(np.c_[RED.ravel(),BLUE.ravel()]))
    boundary=[]
    for i, v in enumerate(np.c_[RED.ravel(),BLUE.ravel()]):
        print(f"{i} : {v.tolist()}")
        res = model.predict(list(v))
        print(res)
        boundary.append(res)
    print(boundary)
    boundary = np.array(boundary)
    plt.figure()
    plt.contourf(RED,BLUE,boundary.reshape(RED.shape),cmap='bwr',alpha=0.3)
    plt.scatter(X_train[:,0],X_train[:,1],c=y_train,cmap='bwr',edgecolor='k')
    plt.xlabel('# red pieces\n[normalized]')
    plt.ylabel('# blue pieces\n[normalized]')
    cbar=plt.colorbar()
    cbar.ax.set_ylabel('Winner')
    plt.title('Centroid Classification')
    plt.show()

    # plot testing data
    if plot_testing == True:
        plt.figure()
        plt.contourf(RED,BLUE,boundary.reshape(RED.shape),cmap='bwr',alpha=0.3)
        testt=plt.scatter(X_test[:,0],X_test[:,1],c=y_test,cmap='bwr',edgecolor='k')
        plt.xlabel('# red pieces\n[normalized]')
        plt.ylabel('# blue pieces\n[normalized]')
        plt.title('Testing Data, n = '+str(X_test.shape[0]))
        cbar=plt.colorbar(testt)
        cbar.ax.set_ylabel('Winner')
        plt.show()
    else:
        pass
    # =============================================================================

    # calculate accuracy with training data set
    predicted=[] # calculate predicted label from training data

    for v in X_test:
        predicted.append(model.predict(v))
    predicted = np.array(predicted)
    acc=np.sum(predicted==y_test)/len(y_test) # calculate accuracy of where prediction is correct vs incorrect relative to actual labels
    print('Training Accuracy =',acc*100,'%\n','='*40,'\n\n')

    return model

def quantum_train(training_data, testing_data, N):
    # Assign training/test data
    X_train = training_data[:,[0,1]] # training features
    y_train = training_data[:,2] # training labels
    X_test = testing_data[:,[0,1]] # testing features
    y_test = testing_data[:,2] # testing labels

    model = QuantumNearestCentroid()

    centroid = model.fit(X_train.tolist(), y_train.flatten().tolist())

    return model

# Uncomment below if you want to run the analysis
#best_quantum = quantum_train_test(training_data=training_data, testing_data=testing_data, N=N, plot_testing=True)

# generate best nearest centroid neighbors
best_quantum = quantum_train(training_data=training_data, testing_data=testing_data, N=N)


compData =  []
compResults = []

def generateCompData(cData = [0,0,0,0,0], qData = [0,0,0,0,0], gData = [0, 0, 0, 0,]): 
    global compData
    global compResults

    ''' Here we take the 2d array gData and transpose it '''
    pgData = pd.DataFrame(gData).T.values.tolist()
    pgData = [0 if x==-1.0 else 1 for x in pgData[2]]

    '''
    These reset the compData and compResults array so there are no errors 
    with the runtime 
    '''
    compData = [
            [0,0,0,0], # Storing what the [Q, C, H, G] say the answer is.
            [0,0,0,0],
            [0,0,0,0],
            [0,0,0,0],
            [0,0,0,0]
    ]
    compResults = [
            [0,0,0], # Storing [QvG, HvG, CvG]
            [0,0,0],
            [0,0,0],
            [0,0,0],
            [0,0,0]
    ]

    ''' 
    The entirety of this function is to setup the enviroment for future functions
    '''
    qCompData = qData # this will be the data from the QML on the 5 boards given
    cCompData = cData # this will be the data from the Classical on the 5 boards given
    gCompData = pgData # this is the ground truth answers

    ''' 
    These take the original qCompData and cCompData arrays and put them 
    together in the compData array for parrsing later in the program 
    '''
    count = 0

    for i in qCompData: 
        compData[count][1] = int(i)
        count += 1
    count = 0
    for i in cCompData:
        compData[count][0] = int(i)
        count += 1
    count = 0
    for i in gCompData:
        compData[count][3] = int(i)
        count += 1

def finalizeCompResults():
    ''' 
    The entirety of this function is to check if the quatum model and 
    human match the conventional model
    '''
    global compData
    global compResults
    count = 0
    for i in compData:
        if i[3] == i[0]: # this is for the C v. G
            compResults[count][1] = 1
        if i[3] == i[1]: # this is for the Q v. G
            compResults[count][0] = 1
        if i[3] == i[2]: # this is for the H v. G
            compResults[count][2] = 1
        count += 1

def showCompResults():
    ''' 
    The entirety of this function is to display to the user which problem they
    or the quantum computer got correct in a visually pleasing way
    '''
    global compResults
    global compData

    data = [
            [0,0,0,0],
            [0,0,0,0],
            [0,0,0,0],
            [0,0,0,0],
            [0,0,0,0]
    ]
    dataColors = [
            ["#a2cffe",0,0,0],
            ["#a2cffe",0,0,0],
            ["#a2cffe",0,0,0],
            ["#a2cffe",0,0,0],
            ["#a2cffe",0,0,0]
    ]

    ''' 
    This next for loop checks the compResults array and sets the values of data
    and dataColors accordingly
    '''
    count = 0
    # print(compResults)
    for i in compResults:
        if i[0] == 1: # Q v. G
            data[count][3] = compData[count][1]
            dataColors[count][3] = "palegreen"
        else:
            dataColors[count][3] = "#cf6275"

        if i[2] == 1: # Human v. G
            data[count][1] = compData[count][2]
            dataColors[count][1] = "palegreen"
        else:
            dataColors[count][1] = "#cf6275"

        if i[1] == 1: # C v. G
            data[count][2] = compData[count][0]
            dataColors[count][2] = "palegreen"
        else:
            dataColors[count][2] = "#cf6275"

        data[count][0] = compData[count][3]
        count += 1
    print(data)
    for i in range(0,5):
        data[i] = ["Red" if x==1 else "Blue" for x in data[i]]
    print(data)
    ''' 
    This section below sums the amount of correct answers from both the human
    and quantum computer
    '''
    Sum = [sum(x) for x in zip(*compResults)]
    winner = ". . . wait what?!?!?! it was a tie! Well, I guess no one"
    if (Sum[0] < Sum[1]) and (Sum[0] < Sum[2]):
        winner = "Human"
    elif (Sum[1] < Sum[0]) and (Sum[1] < Sum[2]):
        winner = "Quantum Computer"
    elif (Sum[2] < Sum[1]) and (Sum[2] < Sum[0]):
        winner = "Classical Computer"
  
    ''' 
    This section below displays the data formated above and who the winner of the 
    game is
    '''
    fig, ax = plt.subplots(1,1)
    column_labels=["Truth", "Human v. Truth", "Classical v. Truth", "Quantum v. Truth"]
    plt.title("Who won, a quantum computer or a human? Turns out the " + winner + " is better!")
    ax.axis('tight')
    ax.axis('off')
    ax.table(cellText=data,
            colLabels = column_labels,
            cellColours = dataColors,
            loc = "center",
            cellLoc='center')
    plt.show()

def showProblem(array):
    ''' 
    This function is used to display the game board for the user to evaluate
    '''
    global compData
    N = 8 # board dimensions: N x N
    board = array
    board = board.reshape(N,N)

    plt.axis('off')
    plt.title("Who has more?")
    plt.imshow(board,cmap='bwr')
    plt.show()

def playersMove(roundNum):
    ''' 
    This function is used to initialize the user to input an answer for the 
    game board
    '''
    global compData
    while True:
        print("Is red or blue more? (b/r): ")
        playerChoice = input("").lower()
        if playerChoice == "r":
            playerChoice = 1
            break
        elif playerChoice == "b":
            playerChoice = 0
            break
    compData[roundNum][2] = playerChoice

clear_output(wait=True)
newGame = True
while newGame:
    # Generate 5 boards for the player to play
    player_data_5, player_images_5, __, __ =generate_data(D=5,N=8,train_split=1)
    cCompData = player_prediction_classical(best_model=best_nCen,player_data=player_data_5)
    # qCompData = player_prediction_quantum(best_model=best_quantum,player_data=player_data_5)
    qCompData=[0,0,0,0,0]
    generateCompData(cCompData, qCompData, player_data_5)
    clear_output(wait=True)

    # Start the game!
    for i in range(0, 5):
        showProblem(player_images_5[i])
        playersMove(i)
        clear_output(wait=True)
    finalizeCompResults()
    showCompResults()
    while True:
        print("Would you like to play again? (y/n): ")
        playerChoice = input("").lower()
        if playerChoice == "y":
            clear_output(wait=True)
            break
        elif playerChoice == "n":
            clear_output(wait=True)
            newGame = False
            break

