import numpy as np
import pickle

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

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


# 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

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]) # Rearranged words
        sentence = words_new >> grammar_dict['ADJ_N_IV'] # Apply grammar on rearranged words
        data_new_psr_dict[sentstr] = sentence.normal_form() # Store sentece with optimum diagram
    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() 

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

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

#*****************************************************
# Translation to quantum circuits
#*****************************************************

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

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}   

#-----------------------------------------
# parametrised part of ansaetze
#-----------------------------------------

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)

#----------------------------------------------------------
# Define parametrised functor to quantum circuits
#----------------------------------------------------------def F(params): 
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})  # send the effect with n.l to same quantum effect
    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
            if arity == 1:
                parshapes.append((p_n,))       
            if arity != 1:
                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))

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
    if labels_dict[sentstr] == '0':
        label = np.array([1,0])
    elif labels_dict[sentstr] == '1':
        label = np.array([0,1])
    psr_labels.append(label)

train_data_psr, test_data_psr, train_labels, test_labels = \
    train_test_split(psr_diagrams, psr_labels, test_size=0.25, random_state=42)

from qiskit import BasicAer, execute, ClassicalRegister
from pytket.extensions.qiskit import tk_to_qiskit
from qiskit.quantum_info import Statevector
#from qiskit import IBMQ
#from qiskit.providers.aer.noise import NoiseModel

backend = BasicAer.get_backend('qasm_simulator')
noise_model = None
coupling_map = None
basis_gates= None

## Build noise model from backend properties
#provider = IBMQ.load_account()
#backend = provider.get_backend('ibmq_16_melbourne')
#noise_model = NoiseModel.from_backend(backend)

## Get coupling map from backend
#coupling_map = backend.configuration().coupling_map

## Get basis gates from noise model
#basis_gates = noise_model.basis_gates

#backend = Aer.get_backend('qasm_simulator')

max_shots = 8192

#with open(f'./noise_model_melbourne.pkl', "rb") as fp:  
    #data = pickle.load(fp)
#noise_model, coupling_map, basis_gates = data[0], data[1], data[2]

from qiskit import Aer

def get_qiskit_results(circ, s_type, retries=10, qasm=False, noise_model=None, 
                       basis_gates=None, coupling_map=None, max_shots=8192):
    """Function to get the Qiskit's results from a DisCoPy circuit.
    
    Args:
        circ = DisCoPy circuit.
        s_type = sentence type for the sentence associated to the 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).
        noise_model = Qiskit's noise model object (Default is None).
        basis_gates = basis gates of the device (Default is None).
        coupling_map = coupling_map of the device (Default is None).
        max_shots = number of shots (Default is 8192).
        
            
    Returns:
        array: results from the postselection measurement.
    """
    
    # Select qubits for postselection according to the sentence type    
    
    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
    
    # From DisCoPy to Qiskit circuit    
    qc = tk_to_qiskit(circ.to_tk())
    if qasm == True:
        #qasm_simulator
        out_reg = ClassicalRegister(1)
        qc.add_register(out_reg)
        qc.measure(post, out_reg[0]) 
        if noise_model is not None:
            backend = Aer.get_backend('qasm_simulator')

            # 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:
            backend = BasicAer.get_backend('qasm_simulator')
            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, retries=retries-1, qasm = qasm,
                                     noise_model=noise_model, coupling_map=coupling_map,
                                     basis_gates=basis_gates, max_shots = max_shots)

        return parse_results(results, eff=zeros)
    
    else:
        # statevector_simulator
        qc.remove_final_measurements()
        state_dict = Statevector(qc).to_dict()
        post0, post1 = '0' * n_qubits, '0' * n_qubits
        post1 = post1[:post] + '1' + post1[post + 1:]
        return [np.abs(state_dict[post0])**2, np.abs(state_dict[post1])**2]
    
def parse_results(results, eff='0'):
    """Parse results of qasm_simulator.
    
    Args:
        results: counts from quantum circuit's measurement.
        eff: digits for splitting the results for post selection.
        
    Returns:
        array: post selected results.
    """
    ks = [k for k in results.keys() if k.split()[1] == eff]
    post_selected_results = {k.split()[0]:v for (k,v) in results.items() if k in ks}
    
    if '0' not in post_selected_results:
        post_selected_results['0'] = 0
    if '1' not in post_selected_results:
        post_selected_results['1'] = 0
        
    total = post_selected_results['0'] + post_selected_results['1']
    
    assert total > 0, 'total was 0 after postselection - try running with more shots'
    
    
    return [post_selected_results['0']/total,  post_selected_results['1']/total]

def process_dataset(dataset):
    """ Process a QNLP dataset.
    
    Args:
        dataset: txt dataset to be processed.
    
    Returns:
        psr_diagrams: parsing diagrams.
        psr_labels: labels.
        sentences: sentences.
        sent_type: dictionary with the sentence type for each sentence.
    """
    with open(dataset) 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
    # Create the previous dictionaries from the given dataset
    adjective_words = [a.name for a in adjectives]
    for sentence in data:
        sentstr = sentence[:-7]
        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'
                else:
                    sent_type[sentstr] = 'tra_' + 'r'
            else:
                sent_type[sentstr] = 'tra'  
    # Rearrange the parsing of the diagrams
    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]) # Rearranged words
            sentence = words_new >> grammar_dict['ADJ_N_IV'] # Apply grammar on rearranged words
            data_new_psr_dict[sentstr] = sentence.normal_form() # Store sentece with optimum diagram
        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() 
    # Store the diagrams, labels and sentences for the classification
    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
        if labels_dict[sentstr] == '0':
            label = np.array([1,0])
        elif labels_dict[sentstr] == '1':
            label = np.array([0,1])
        psr_labels.append(label)
        
    return psr_diagrams, psr_labels, sentences, sent_type

from qiskit.aqua.components.optimizers import SPSA
import time

def QNLP_optimize(dataset, initial_point, n_iterations, cost_function, qasm_bool=False,
                  noise_model=None, coupling_map=None, basis_gates=None, max_shots = 8192):
    """ Optimize a QNLP dataset.
    
    Args:
        dataset = dataset for the optimization.
        initial_point = initial value of the parameters.
        niter = number of iterations for SPSA optimizer.
        cost_function = cost function for the optimization
        qasm_bool: if True qasm_simulator is used, if not
            statevector_simulator (Default is False).
    Returns:
        list: Qiskit's SPSA results.
        cost_save: cost of each evaluation of the cost function.
        params_save: parameters of each evaluation of the cost function.
    """
    # Obtain the data from the dataset
    train_data_psr, train_labels, train_sent, sent_type = process_dataset(dataset)

    # Qiskit's SPSA parameters
    num_vars = len(initial_point)
    bounds = [[0.0, 1.0] for _ in range(len(initial_point))]
    
    # Initialize cost function global parameters
    cost_save = []
    params_save = []
    i_cost = []
    # Choose cost function
    if cost_function == 'cost_train':
        def get_cost_train_error(unshaped_params):
            """ Cost function using training error.

            Args:
                unshaped_params: parameters for the optimization.

            Returns:
                float: cost.
            """
            # Print iterations
            i_cost.append(1)
            if len(i_cost) % 2 == 0:
                print('Iteration', len(i_cost)/2)

            # Create functor to construct DisCoPy circuits
            func = F(reshape_params(unshaped_params, par_shapes))
            circs = [func(diag) for diag in train_data_psr]
            # Get Qiskit's results
            results = np.array([get_qiskit_results(circs[i], sent_type[train_sent[i]], qasm = qasm_bool,
                                     noise_model=noise_model, coupling_map=coupling_map,
                                     basis_gates=basis_gates, max_shots=max_shots) for i in range(len(train_sent))])
            #Tweak results
            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]
            # Compute cost
            error = 0
            for i, r in enumerate(pred_labels_distrs):
                error += int(np.argmax(r) != np.argmax(train_labels[i]))
            # Print cost
            print('Cost', error/len(pred_labels_distrs))
            # Save cost and parameters
            cost_save.append(error/len(pred_labels_distrs))
            params_save.append(unshaped_params)
            return error/len(pred_labels_distrs)
        
        get_cost = get_cost_train_error
    
    elif cost_function == 'cost_entropy':

        def get_cost_cross_entropy(unshaped_params):
            """ Cost function using cross entropy.

            Args:
                unshaped_params: parameters for the optimization.

            Returns:
                float: cost.
            """
            # Print iterations
            i_cost.append(1)
            if len(i_cost) % 2 == 0:
                print('Iteration', len(i_cost)/2)

            # Create functor to construct DisCoPy circuits
            func = F(reshape_params(unshaped_params, par_shapes))
            circs = [func(diag) for diag in train_data_psr]
            # Get Qiskit's results
            results = np.array([get_qiskit_results(circs[i], sent_type[train_sent[i]], qasm = qasm_bool,
                                     noise_model=noise_model, coupling_map=coupling_map,
                                     basis_gates=basis_gates, max_shots=max_shots) for i in range(len(train_sent))])#Tweak results
            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]
            # Compute cost
            cross_entropies = np.array([np.sum(train_labels[s] * np.log2(pred_labels_distrs[s])) for s in range(len(train_labels))])
            # Print cost
            print('Cost', -1 / len(train_data_psr) * np.sum(cross_entropies))
            # Save cost and parameters
            cost_save.append(-1 / len(train_data_psr) * np.sum(cross_entropies))
            params_save.append(unshaped_params)
            return -1 / len(train_data_psr) * np.sum(cross_entropies)
            
        get_cost = get_cost_cross_entropy
    else:
        raise NameError('No valid cost function name.')
    # Optimization
    t1 = time.time()
    opt = SPSA(maxiter=n_iterations)
    ps = opt.optimize(num_vars, get_cost, initial_point=initial_point, variable_bounds=bounds)
    print("Final time", time.time() - t1)
    if qasm_bool == True:
        simul = 'qasm'
    else:
        simul = 'statevector'
    pickle.dump([ps, cost_save, params_save], 
                open(f"sentiment_analysis_1vs1_{dataset[16:-4]}_{str(cost_function)}_niter{n_iterations}_{simul}.pkl", "wb"))
    return ps, cost_save, params_save

datasets = ['train_happy_vs_sad.txt', 'train_happy_vs_angry.txt', 'train_happy_vs_scared.txt',
            'train_sad_vs_angry.txt', 'train_sad_vs_scared.txt', 'train_angry_vs_scared.txt'] 
n_iterations = 1
cost_function = 'cost_train'
initial_point = randparams(par_shapes)
for dataset in datasets:
    dataset = './1vs1_datasets/' + dataset
    qasm_bool = False
    results = QNLP_optimize(dataset,initial_point, n_iterations, cost_function, qasm_bool=qasm_bool,
                            noise_model=noise_model, coupling_map=coupling_map,
                            basis_gates=basis_gates, max_shots=max_shots)

import qiskit
qiskit.__qiskit_version__