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

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

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

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

from discopy import Diagram
from discopy.grammar import draw

# Store the grammatical structures of each sentence type 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 = {}

#Creates all possible sentences

# 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}


# 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]

        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
            

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

# 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
from discopy.quantum.gates import Rx, Rz

import numpy as np

q_s = 1 # number of qubits for type s (sentence)
q_n = 1 # number of qubits for type n (noun)
depth = 1 # depth of the IQPansatz
p_n = 3 # number of parameters for the single qubit iqp ansatz

# 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 single_qubit_iqp_ansatz(params):
    if len(params) == 1:
        return Rx(params[0])  
    if len(params) == 2:
        return Rx(params[0]) >> Rz(params[1])
    if len(params) == 3:
        return IQPansatz(1, params)     
    
def ansatz_state(state, params):
    arity = sum(ob[Ty(factor.name)] for factor in state.cod) 
    if arity == 1:
        return Ket(0) >> single_qubit_iqp_ansatz(params)
    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)
    if arity == 1:
        return single_qubit_iqp_ansatz(params) >> Bra(0)
    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

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

#****************************************
# 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 = []

#create lists of sentences and corresponding diagram and labels
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
    if labels_dict[sentstr] == '0':
        label = np.array([1,0,0,0])
    elif labels_dict[sentstr] == '1':
        label = np.array([0,1,0,0])
    elif labels_dict[sentstr] == '2':
        label = np.array([0,0,1,0])
    elif labels_dict[sentstr] == '3':
        label = np.array([0,0,0,1])
    psr_labels.append(label)
    
    
#Splits the sentences and corresponding labels into train and test sets
train_sent, test_sent, train_labels, test_labels = \
    train_test_split(sentences, psr_labels, test_size=0.25, random_state=42)

#Finds the discopy diagrams for the sentences in train and test
train_data_psr = [psr_diagrams_dict[sent] for sent in train_sent]
test_data_psr = [psr_diagrams_dict[sent] for sent in test_sent]

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

#intransitive sentence
func(psr_diagrams_dict['morose woman cries.']).draw()

#basic transitive sentence
func(psr_diagrams_dict['man grieves woman.']).draw()

#transitive sentence, left adjective
func(psr_diagrams_dict['morose man grieves woman.']).draw()

#transitive sentence, right adjective
func(psr_diagrams_dict['man entertains joyful woman.']).draw()

def get_qiskit_results(circ, sent, retries=10, qasm=False):
    """Function to get the Qiskit's results from a DisCoPy circuit.
    
    Args:
        circ: DisCoPy circuit.
        sent: English sentence corresponding to circuit
        retries: number of retries before assigning a random result
            if the postselection leads to zero values (Default is 10).
        qasm: if True qasm_simulator is used, if not statevector_simulator
            (Default is False).
            
    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 = 8
        posts = [4,5]
    elif s_type == 'tra':
        n_qubits = 6
        posts = [2,3]
    elif s_type == 'tra_l':
        n_qubits = 10
        posts = [4,5]
    else:
        assert s_type == 'tra_r'
        n_qubits = 10
        posts = [2,3]
    
    
    #the diagrams are initially in discopy, so we first convert them to tket and then to qiskit
    qc = tk_to_qiskit(circ.to_tk())
    
    if qasm == True:
        backend = Aer.get_backend('qasm_simulator')
        #qasm_simulator
        out_reg = ClassicalRegister(2)
        qc.add_register(out_reg)
        qc.measure(posts, out_reg) 
        if noise_model is not None:
            # Include noise model
            results = execute(qc, backend, shots=max_shots, noise_model=noise_model, 
                              coupling_map=coupling_map,basis_gates=basis_gates).result().get_counts()
        else:
            results = execute(qc, backend, shots=max_shots).result().get_counts()

        zeros = '0' * (n_qubits - 1)

        if '0 ' + zeros not in results and '1 ' + zeros not in results:
            if retries == 0:
                return np.random.rand(2)
            return get_qiskit_results(circ, s_type, transpile, retries=retries-1, qasm = qasm_bool)

        return parse_results(results, eff=zeros)
    
    
    else:

        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 4 possible values
        #for the sentence qubits, which we store in a list here
        selections = ['0' * posts[0] + val + '0' * (n_qubits - posts[1] - 1) for val in ['00', '01', '10', '11']]

        return [np.abs(state_dict[sel])**2 for sel in selections] #returns the amplitude of each of these four states
    

from qiskit.algorithms.optimizers import SPSA
import time

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

begin_time = time.time()

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


ps = {}

#Chooses random parameters to begin with
rand_unshaped_pars = randparams(par_shapes)
num_vars = len(rand_unshaped_pars)

#create bounds between 0 and 1 for every variable
bounds = [[0.0, 1.0] for _ in range(len(rand_unshaped_pars))]

n_iterations = 1250

a = time.time()

#creates a qiskit SPSA optimizer
opt = SPSA(maxiter=n_iterations)

print("\n\ncalibrating")
opt.calibrate(cost_ce, initial_point=rand_unshaped_pars) #calibrates it with the cost function and starting parameters
print("training")

global_costs = []
begin_time = time.time()

ps = opt.optimize(num_vars, cost_ce, initial_point=rand_unshaped_pars, variable_bounds=bounds)
print("took", time.time() - a)
pickle.dump([ps, global_costs], open("sentiment_analysis_qiskit_ce_1250.p", "wb"))

func = F(reshape_params(ps[0], par_shapes))
    
circs = [func(diag) for diag in test_data_psr]
    
results = np.array([get_qiskit_results(circs[i], test_sent[i]) for i in range(len(test_sent))])

    
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]

from sklearn.metrics import confusion_matrix
import pandas as pd
import seaborn as sns

c_matrix = confusion_matrix([np.argmax(t) for t in test_labels], [np.argmax(r) for r in pred_labels_distrs])

df_cm = pd.DataFrame(c_matrix, ['happy','sad','angry','scared'], ['happy','sad','angry','scared'])
df_cm.index.name = 'Actual'
df_cm.columns.name = 'Predicted'
sns.heatmap(df_cm, annot=True, annot_kws={"size": 16},cmap='viridis',
                xticklabels=True, yticklabels=True)

print("test accuracy =", sum([np.argmax(test_labels[i]) == np.argmax(pred_labels_distrs[i]) for i in range(len(test_labels))])/len(test_labels))

costs = global_costs[51:]
n = ps[-1]
costs = [(costs[j*2] + costs[j*2 + 1])/2 for j in range(n//2)]

import matplotlib.pyplot as plt


plt.xlabel('Iterations')
plt.ylabel('Cross Entropy')
plt.title("2-Qubit Multi-class Classification")
plt.plot(costs)


