import numpy as np
import pickle

import qiskit
qiskit.__qiskit_version__

from discopy import Ty, Word # Import the classes for the type of word and for the word

s, n = Ty('s'), Ty('n') # Define the types s and n
nphr, adj, tv, iv, vphr = Ty('NP'), Ty('ADJ'), Ty('TV'), Ty('IV'), Ty('VP')

# Define the words (notice that we include both meaning and grammar)

# nouns
man, woman, kid = Word('man', n), Word('woman', n), Word('kid', n)
# adjectives
morose, irascible = Word('morose', n @ n.l), Word('irascible', n @ n.l)
frightened, cheerful = Word('frightened', n @ n.l), Word('cheerful', n @ n.l)
gloomy, furious = Word('gloomy', n @ n.l), Word('furious', n @ n.l)
terrified, joyful = Word('terrified', n @ n.l), Word('joyful', n @ n.l)
downcast, miserable = Word('downcast', n @ n.l), Word('miserable', n @ n.l)
old, young = Word('old', n @ n.l), Word('young', n @ n.l)
# Intransitive verbs
cries, shouts = Word('cries', n.r @ s), Word('shouts', n.r @ s)
laughs, snaps = Word('laughs', n.r @ s), Word('snaps', n.r @ s)
# Transitive verbs
grieves, startles = Word('grieves', n.r @ s @ n.l), Word('startles', n.r @ s @ n.l)
entertains, irritates = Word('entertains', n.r @ s @ n.l), Word('irritates', n.r @ s @ n.l)

nouns = [man, woman, kid]
adjectives = [morose, irascible, frightened, cheerful, gloomy, furious, terrified, joyful, downcast, miserable, old, young]
int_verbs = [cries, shouts, laughs, snaps]
t_verbs = [grieves, startles, entertains, irritates]

vocab = nouns + int_verbs + t_verbs + adjectives

from discopy import Cup, Id, pregroup

grammar =  Id(n) @ Cup(n.l, n) @ Id(n.r @ s) >> Cup(n, n.r) @ Id(s) >> Id(s)
sentence = joyful @ woman @ laughs >> grammar
pregroup.draw(sentence)

grammar =  Cup(n, n.r) @ Id(s) @ Cup(n.l, n)
sentence = woman @ grieves @ kid >> grammar
pregroup.draw(sentence)

from discopy import Diagram
from discopy.grammar import draw

# Store the grammatical structures in a dictionary 

grammar_dict = {
    'ADJ_N_IV' : Id(n) @ Cup(n.l, n) @ Id(n.r @ s) >> Cup(n, n.r) @ Id(s) >> Id(s),
    'N_TV_N': Cup(n, n.r) @ Id(s) @ Cup(n.l, n),
    'ADJ_N_TV_N': Id(n) @ Cup(n.l, n) @ Id(n.r @ s) @ Cup(n.l, n) >> Cup(n, n.r) @ Id(s) >> Id(s),
    'N_TV_ADJ_N': Cup(n, n.r) @ Id(s) @ Cup(n.l, n) @ Cup(n.l, n) >> Id(s)}
                                                                          

# Create parsing (grammatical analysis) dictionary where the grammatical sentences
# are the keys and the associated values are the diagrams (words + grammar)

data_psr = {}

# Intransitive sentences
parsing_int = {"{} {} {}.".format(adj, noun, int_verb): adj @ noun @ int_verb >> grammar_dict['ADJ_N_IV']
            for adj in adjectives for noun in nouns for int_verb in int_verbs}
sentences_int = list(parsing_int.keys())
for sentence in sentences_int:
    diagram = parsing_int[sentence]
    data_psr[sentence] = parsing_int[sentence]

# Transitive sentences (without adjective)
parsing_tra = {"{} {} {}.".format(subj, t_verb, obj):  subj @ t_verb @ obj >> grammar_dict['N_TV_N']
            for subj in nouns for t_verb in t_verbs for obj in nouns}


#Due to time limitations, we are exlcuding transitive sentences with adjectives in this implementation
"""
# Transitive sentences (with adjective)
parsing_tra_ladj = {"{} {} {} {}.".format(adj, subj, t_verb, obj):  adj @ subj @ t_verb @ obj >> grammar_dict['ADJ_N_TV_N']
            for adj in adjectives for subj in nouns for t_verb in t_verbs for obj in nouns}
parsing_tra_radj = {"{} {} {} {}.".format(subj, t_verb, adj, obj):  subj @ t_verb @ adj @ obj >> grammar_dict['N_TV_ADJ_N']
            for subj in nouns for t_verb in t_verbs for adj in adjectives for obj in nouns}

parsing_tra.update(parsing_tra_ladj) #merges transitive adjectives into original dict
parsing_tra.update(parsing_tra_radj)
"""
sentences_tra = list(parsing_tra.keys())
for sentence in sentences_tra:
    diagram = parsing_tra[sentence]
    data_psr[sentence] = parsing_tra[sentence]

with open('../sentiment_analysis_dataset.txt') as f:
    data = f.readlines()

labels_dict = {} # Dictionary with the labels for each sentence
data_psr_dict = {} # Dictionary with the parsing for each sentence
sent_type = {} # Dictionary with the sentence type for each sentence


adjective_words = [a.name for a in adjectives]
for sentence in data:
    sentstr = sentence[:-7] #seperates the sentence string from the data at the end
    if sentence[-6:-3] == 'int':
        diagram = parsing_int[sentstr]
        data_psr_dict[sentstr] = diagram
        labels_dict[sentstr] = sentence[-2]
        sent_type[sentstr] = 'int'
    elif sentence[-6:-3] == 'tra':
        diagram = parsing_tra[sentstr]
        data_psr_dict[sentstr] = diagram
        labels_dict[sentstr] = sentence[-2]

        #ignoring transitive adjective sentences in this run
        """if len(sentstr.split()) == 4:
            if sentstr.split()[0] in adjective_words:
                sent_type[sentstr] = 'tra_' + 'l' #adjective on the left
            else:
                sent_type[sentstr] = 'tra_' + 'r' #adjective on the right
        else:
            sent_type[sentstr] = 'tra' #the simple transitive verb sentence type
        """
        sent_type[sentstr] = 'tra'
            

c0 = 0
c1 = 0
c2 = 0
c3 = 0
for value in labels_dict.values():
    if value == '0':
        c0 +=1
    elif value == '1':
        c1 += 1
    elif value == '2':
        c2 += 1
    elif value == '3':
        c3 += 1
print('Number of elements for each sentiment')
print('Happy: ', c0)
print('Sad: ', c1)
print('Angry: ', c2)
print('Scared: ', c3)
print('Total', len(data_psr_dict))

from discopy import Cap, Box

woman2 = Box('woman', n.l, Ty())
words_new = (Cap(n, n.l)) >> (joyful @ Id(n) @ woman2 @ laughs)
sentence = words_new >> grammar_dict['ADJ_N_IV']
sentence = sentence.normal_form()

sentence.draw()

woman2 = Box('woman', n.r, Ty())
kid2 = Box('kid', n.l, Ty())
words_new = (Cap(n.r, n) @ Cap(n, n.l)) >> (woman2 @ Id(n) @ grieves @ Id(n) @ kid2)
sentence = words_new >> grammar_dict['N_TV_N']
sentence = sentence.normal_form()
sentence.draw()

data_psr['morose man grieves woman.'].draw()

# Code of this cell addapted from https://github.com/CQCL/qnlp_lorenz_etal_2021_resources
#iterates through all sentences and reshapes them as above according to their type

data_new_psr_dict = {}
for sentstr in data_psr_dict.keys(): 
    num_words = len(sentstr.split(' '))
    words = data_psr_dict[sentstr].boxes[:num_words]
    if sent_type[sentstr] == 'int':
        noun = Box(words[1].name, n.l, Ty())
        words_new = (Cap(n, n.l)) >> (words[0] @ Id(n) @ noun @ words[2])
        sentence = words_new >> grammar_dict['ADJ_N_IV']
        data_new_psr_dict[sentstr] = sentence.normal_form()
    elif 'tra' in sent_type[sentstr]:
        if sent_type[sentstr] == 'tra':
            noun1 = Box(words[0].name, n.r, Ty())
            noun2 = Box(words[2].name, n.l, Ty())
            words_new = (Cap(n.r, n) @ Cap(n, n.l)) >> (noun1 @ Id(n) @ words[1] @ Id(n) @ noun2)
            sentence = words_new >> grammar_dict['N_TV_N']
            data_new_psr_dict[sentstr] = sentence.normal_form()
    """
        elif sent_type[sentstr] == 'tra_l': #adjective at beginning
            noun1 = Box(words[1].name, n.l, Ty())
            noun2 = Box(words[3].name, n.l, Ty())
            words_new = (Cap(n, n.l) @ Cap(n, n.l)) >> (words[0] @ Id(n) @ noun1 @ words[2] @ Id(n) @ noun2)
            sentence = words_new >> grammar_dict['ADJ_N_TV_N']
            data_new_psr_dict[sentstr] = sentence.normal_form()
        elif sent_type[sentstr] == 'tra_r': #adjective on second noun
            noun1 = Box(words[0].name, n.r, Ty())
            noun2 = Box(words[3].name, n.l, Ty())
            words_new = (Cap(n.r, n) @ Cap(n, n.l)) >> (noun1 @ Id(n) @ words[1] @ words[2] @ Id(n) @ noun2)
            sentence = words_new >> grammar_dict['N_TV_ADJ_N']
            data_new_psr_dict[sentstr] = sentence.normal_form() 
    """

vocab_psr = []
for word in vocab:
    if word.cod == Ty('n'):
        vocab_psr.append(Box(word.name, n.r, Ty()))   # n.l case is dealt with in definition of quantum functor
    else:
        vocab_psr.append(word)

# Code of this cell from https://github.com/CQCL/qnlp_lorenz_etal_2021_resources


from discopy.quantum import Ket, IQPansatz, Bra, qubit
from discopy.quantum.circuit import Functor, Id
from discopy.quantum.circuit import Circuit
from functools import reduce, partial

import numpy as np

q_s = 2 # number of qubits for type s (sentence)
q_n = 2 # number of qubits for type n (noun)
depth = 1 # depth of the IQPansatz

# Define the dimensions of the objects of the circuit functor
ob = {s: q_s, n: q_n}                          
ob_cqmap = {s: qubit ** q_s, n: qubit ** q_n}   

# Define the ansätze for states and effects
def ansatz_state(state, params):
    arity = sum(ob[Ty(factor.name)] for factor in state.cod) 
    return Ket(*tuple([0 for i in range(arity)])) >> IQPansatz(arity, params)
    
def ansatz_effect(effect, params): 
    arity = sum(ob[Ty(factor.name)] for factor in effect.dom)
    return IQPansatz(arity, params) >> Bra(*tuple([0 for i in range(arity)]))
       
def ansatz(box,params):
    dom_type = box.dom
    cod_type = box.cod
    if len(dom_type) == 0 and len(cod_type) != 0:
        return ansatz_state(box, params)
    if len(dom_type) != 0 and len(cod_type) == 0:
        return ansatz_effect(box, params)

# Construct the circuit functor
def F(params): 
    ar = dict()
    for i in range(len(vocab_psr)):
        pgbox = vocab_psr[i]
        qbox = ansatz(vocab_psr[i], params[i])
        ar.update({pgbox: qbox})
        if pgbox.cod == Ty():
            ar.update({Box(pgbox.name, n.l, Ty()): qbox})
    return Functor(ob_cqmap, ar)

# Code of this cell from https://github.com/CQCL/qnlp_lorenz_etal_2021_resources

#*****************************************************
# Functions to deal with the parametrisation
#*****************************************************

def paramshapes(vocab_psr):
    parshapes = []    
    for box in vocab_psr:
        dom_type = box.dom
        cod_type = box.cod
        dom_arity = sum(ob[Ty(factor.name)] for factor in box.dom)
        cod_arity = sum(ob[Ty(factor.name)] for factor in box.cod)
        if dom_arity == 0 or cod_arity == 0:  # states and effects
            arity = max(dom_arity, cod_arity)
            assert arity != 0
            parshapes.append((depth, arity-1))
    return parshapes

def randparams(par_shapes):
    params = np.array([]) 
    for i in range(len(par_shapes)):
        params = np.concatenate((params, np.ravel(np.random.rand(*par_shapes[i])))) # np.ravel flattens an array
    return params 

def reshape_params(unshaped_pars, par_shapes):
    pars_reshaped = [[] for ii in range(len(par_shapes))]
    shift = 0
    for ss, s in enumerate(par_shapes):
        idx0 = 0 + shift
        if len(s) == 1:
            idx1 = s[0] + shift
        elif len(s) == 2:
            idx1 = s[0] * s[1] + shift
        pars_reshaped[ss] = np.reshape(unshaped_pars[idx0:idx1], s)
        if len(s) == 1:
            shift += s[0]
        elif len(s) == 2:
            shift += s[0] * s[1]
    return pars_reshaped

#****************************************
# Parameters of the current model
#****************************************

par_shapes = paramshapes(vocab_psr)
rand_unshaped_pars = randparams(par_shapes)
rand_shaped_pars = reshape_params(rand_unshaped_pars, par_shapes)

print('Number of parameters:    ', len(rand_unshaped_pars))

# Print the quantum circuit for each sentence

func = F(rand_shaped_pars)

for sentstr in data_new_psr_dict:
    print(sentstr)
    print(data_new_psr_dict[sentstr])
    #func(data_new_psr_dict[sentstr]).draw(draw_box_labels=True, figsize=(5, 5))

from sklearn.model_selection import train_test_split

psr_diagrams = []
psr_diagrams_dict = {}
psr_labels = []
sentences = []

for sentstr in data_new_psr_dict.keys():
    sentences.append(sentstr)
    diagram = data_new_psr_dict[sentstr]
    psr_diagrams.append(diagram)
    psr_diagrams_dict[sentstr] = diagram
    label = int(labels_dict[sentstr])
    psr_labels.append(label)

#create lists of sentences and corresponding diagram and labels
train_data_psr, test_data_psr, orig_train_labels, orig_test_labels = \
    train_test_split(psr_diagrams, psr_labels, test_size=0.25, random_state=42)
train_sent, test_sent, train_labels_sent, test_labels_sent = \
    train_test_split(sentences, psr_labels, test_size=0.25, random_state=42)

def num_to_vec(x):
    if x:
        return [0,1] #True
    else:
        return [1,0]

data_sets_labels = {}
for i, cla in enumerate(['happy', 'sad', 'angry', 'scared']):
    data_sets_labels[cla] = {}
    
    tr_labels = [num_to_vec(i == t) for t in orig_train_labels]
    data_sets_labels[cla]['train'] = tr_labels
    
    te_labels = [num_to_vec(i == t) for t in orig_test_labels]
    data_sets_labels[cla]['test'] = te_labels

from qiskit import BasicAer, execute, ClassicalRegister
from pytket.extensions.qiskit import tk_to_qiskit
from qiskit.quantum_info import Statevector

func(data_new_psr_dict['morose woman cries.']).draw()

func(data_new_psr_dict['woman grieves man.']).draw()

def get_qiskit_results(circ, sent, retries=10):
    """Function to get the Qiskit's results from a DisCoPy circuit.
    
    Args:
        circ: DisCoPy circuit.
        sent: English sentence for circ
            
    Returns:
        array: results from the postselection measurement.
    """
    
    #Finds the sentence's type to determine the number of qubits and position of the sentence qubits    
    s_type = sent_type[sent]
    
    if s_type == 'int':
        n_qubits = 4
        post = 2
    elif s_type == 'tra':
        n_qubits = 3
        post = 1
    elif s_type == 'tra_l':
        n_qubits = 5
        post = 2
    else:
        assert s_type == 'tra_r'
        n_qubits = 5
        post = 1
    
    #the diagrams are initially in discopy, so we first convert them to tket and then to qiskit    
    qc = tk_to_qiskit(circ.to_tk())
    

    
    qc.remove_final_measurements()
    
    state_dict = Statevector(qc).to_dict()
    
    #The only values we care about are the ones when the non-sentence qubits are all zero.
    #This leaves only two possible states - the sentence qubit is 0 or 1 (and the rest are all 0)
    selections = ['0' * post + val + '0' * (n_qubits - post - 1) for val in ['0', '1']]
    
    
    return [np.abs(state_dict[sel])**2 for sel in selections] #returns the amplitude for the two states
      

from qiskit.algorithms.optimizers import SPSA
import time


global_costs = [] #global variable to keep track of cost during training


def cost_ce(unshaped_params): #finds cross entropy cost
    
    #creates a functor which turns a sentence diagram into a circuit parameterized by unshaped_params
    func = F(reshape_params(unshaped_params, par_shapes)) 
    
    #Applies this functor to every training diagram
    circs = [func(diag) for diag in train_data_psr]
    
    #Finds the post-selected qiskit results for each circuit
    results = np.array([get_qiskit_results(circs[i], train_sent[i]) for i in range(len(train_sent))])

    #Turns each result into a probability distribution
    results_tweaked = [np.abs(np.array(res) - 1e-9) for res in results] 
    pred_labels_distrs = [res.flatten() / np.sum(res) for res in results_tweaked]
    
    
    #finds the cross entropy
    cross_entropies = np.array([np.sum(train_labels[s] * np.log2(pred_labels_distrs[s])) for s in range(len(train_labels))])
    cost = -1 / len(train_data_psr) * np.sum(cross_entropies)
    
    #stores the cost in a global list so we can track progress
    global_costs.append(cost)
    
    if len(global_costs) % 100 == 0: #since SPSA calls the cost function twice per iteration, this prints every 50 iterations
        print("at iteration", len(global_costs)/2, time.time() - begin_time)
       
    
    return cost



rand_unshaped_pars = randparams(par_shapes)
num_vars = len(rand_unshaped_pars)

bounds = [[0.0, 1.0] for _ in range(len(rand_unshaped_pars))]

ps = {}


n_iterations = 1000

#creates a classifier 1vall classifier for each sentiment
for cla in ['happy', 'sad', 'angry', 'scared']:
    
    #chooses random starting parameters
    rand_unshaped_pars = randparams(par_shapes)


    #finds the labels for the current classifier
    train_labels = data_sets_labels[cla]['train']
    test_labels = data_sets_labels[cla]['test']




    a = time.time()

    #creates a qiskit optimizer
    opt = SPSA(maxiter=n_iterations)
    
    print("\n\ncalibrating", cla)
    opt.calibrate(cost_ce, initial_point=rand_unshaped_pars) #calibrates according to cost function
    print("training")
    global_costs = []
    o = opt.optimize(num_vars, cost_ce, initial_point=rand_unshaped_pars, variable_bounds=bounds) #performs optimization
    ps[cla] = [o, global_costs.copy()] #saves parameters and cost history
    print("took", time.time() - a)


final_train_predictions = [(0,'') for _ in range(len(orig_train_labels))] # list of tuples of the form (probability, class)
final_test_predictions = [(0,'') for _ in range(len(orig_test_labels))]

test_predictions_dict = {}

individual_train_predictions = {}
individual_test_predictions = {}
classes = ['happy', 'sad', 'angry', 'scared']
for j, cla in enumerate(classes):
    
    params = ps[cla][0][0]
    func = F(reshape_params(params, par_shapes))
    
    
    train_data = train_data_psr
    #train_data = data_sets[cla]['train']
    
    train_labels = data_sets_labels[cla]['train']
    test_labels = data_sets_labels[cla]['test']
    
    
    final_train_circuits = [func(diag) for diag in train_data_psr]
    train_results = [get_qiskit_results(final_train_circuits[i], train_sent[i]) for i in range(len(train_sent))]

    final_test_circuits = [func(diag) for diag in test_data_psr]
    test_results = [get_qiskit_results(final_test_circuits[i], test_sent[i]) for i in range(len(test_sent))]

    train_results_tweaked = [np.abs(np.array(res) - 1e-9) for res in train_results]
    test_results_tweaked = [np.abs(np.array(res) - 1e-9) for res in test_results]
    pred_train_results = [res.flatten() / np.sum(res) for res in train_results_tweaked]
    pred_test_results = [res.flatten() / np.sum(res) for res in test_results_tweaked]
    
    
    test_predictions_dict[cla] = ([np.argmax(r) for r in pred_test_results], test_labels)
    
    
    
    #Since each tuple (p, j) in these lists represents a probability that sentence i is in class j, this update keeps 
    #track of the most likely prediction so far
    final_train_predictions = [max(final_train_predictions[i], (pred_train_results[i][1], j)) for i in range(len(pred_train_results))]
    final_test_predictions = [max(final_test_predictions[i], (pred_test_results[i][1], j)) for i in range(len(pred_test_results))]
    
    individual_train_predictions[cla] = [np.argmax(pred) for pred in pred_train_results]
    individual_test_predictions[cla] = [np.argmax(pred) for pred in pred_test_results]

#individual accuracy

for cla in classes:
    train_pred = individual_train_predictions[cla]
    train_actual = data_sets_labels[cla]['train']
    
    acc = [train_pred[i] == np.argmax(train_actual[i]) for i in range(len(train_pred))]
    print("\n\nfor", cla, ": train accuracy =", sum(acc)/len(acc))
    

    test_pred = individual_test_predictions[cla]
    test_actual = data_sets_labels[cla]['test']
    
    acc = [test_pred[i] == np.argmax(test_actual[i]) for i in range(len(test_pred))]
    print("for", cla, ": test accuracy =", sum(acc)/len(acc))
    
    

from sklearn.metrics import confusion_matrix
for cla in classes:
    
    print("\n\n for", cla, "\n", confusion_matrix([np.argmax(t) for t in test_predictions_dict[cla][1]], test_predictions_dict[cla][0])) #true, pred
    
confusion_matrix(orig_test_labels, [r[1] for r in final_test_predictions])

correct = 0
for i, res in enumerate(final_train_predictions):
    pred_result = res[1]
    train_result = orig_train_labels[i]
    #print(train_sent[i])
    #print(res.flatten())
    #print(pred_result,train_result)
    if train_result == pred_result:
        correct += 1
    #print(f'Result: {train_array, train_result}, Predicted result: {res, pred_result}')
print('Correct overall predictions (train):',correct/len(train_results_tweaked))

correct = 0
for i, res in enumerate(final_test_predictions):
    pred_result = res[1]
    test_result = orig_test_labels[i]
    if test_result == pred_result:
        correct += 1
print('Correct overall predictions (test):',correct/len(test_results_tweaked))

pickle.dump(ps, open("results/sentiment_analysis_qiskit_1vall.p", "wb"))

fig, axs = plt.subplots(2,2, figsize=(15,10))
import matplotlib.pyplot as plt

colours = ['blue', 'red', 'orange', 'green']

for i in range(4):
    costs = ps[classes[i]][1][51:] #excluding the first 51 values since they are used for calibration
    
    n = ps[cla][0][-1]
    costs = [(costs[j*2] + costs[j*2 + 1])/2 for j in range(n//2)] #averages over xplus and xminus values
    
    axs[i//2, i%2].plot(ps[classes[i]][1][51:], color=colours[i]) 
    axs[i//2, i%2].set_title(classes[i])

    axs[i//2, i%2].set_xlabel('iterations')
    axs[i//2, i%2].set_ylabel('cost')
    axs[i//2, i%2].set_ylim(0,2.5)


