# %%
import pandas as pd

from .basic_queries import basicQueries


# Executes query sequences to recover contexts with co-occurrences according to user input
# Returns Pandas dataframes
class cooccorrenze(basicQueries):
    
    def __init__(self, dataConfig):
        super().__init__(dataConfig)

#%% funzione ricerca per cooccorrenze. 
    # Ha in input un array del tipo [forma/lemma_cercati, tipo_ricerca, ricerca_espansa, iniziale_raddoppiata].
    # l'attributo tipo_ricerca definisce il tipo di ricerca in input (0 per forme, 1 per lemmi, 2 per lemmi con opzione "mostra occorrenze non lemmatizzate").
    # Permette di definire l'intervallo di ricerca (in numero di parole), la possibilità di cercare soltanto all'interno dello stesso periodo (0/1) e/o di cercare le occorrenze in modo ordinato (0/1)
    
    def ricerca_cooccorrenze(self, listaricerche, intervallo, periodo, ordinate):
        listatesti = pd.DataFrame()
        cod = 1
        if listaricerche[0][1] == 0:
            ricerca = self.sendBasicQuery(listaricerche[0][0], 'forma', listaricerche[0][2], listaricerche[0][3], pandas=True)
            listatesti = self.findtexts(0, ricerca)
        elif listaricerche[0][1] == 1:
            ricerca = self.sendBasicQuery(listaricerche[0][0], 'lemma', listaricerche[0][2], listaricerche[0][3], pandas=True)
            listatesti = self.findtexts(1, ricerca)
        elif listaricerche[0][1] == 2:
            ricerca = self.sendBasicQuery(listaricerche[0][0], 'lemma', listaricerche[0][2], listaricerche[0][3], pandas=True)
            listatesti = self.findtexts(2, ricerca)

        if listatesti.empty:
            return []

        for ricerca, tipo, espansa, raddoppiata in listaricerche[1:]:
            if tipo == 0:
                search = self.sendBasicQuery(ricerca, 'forma', espansa, raddoppiata, pandas=True)
            elif tipo == 1:
                search = self.sendBasicQuery(ricerca, 'lemma', espansa, raddoppiata, pandas=True)
            elif tipo == 2:
                search = self.sendBasicQuery(ricerca, 'lemma', espansa, raddoppiata, pandas=True)

            textlist = self.findtexts(tipo, search)

            listatesti = listatesti.merge(textlist, on='ntx', suffixes=('', f'_{cod}'))
            cond1 = listatesti['numperiod'] == listatesti[f'numperiod_{cod}'] if periodo == 1 else True
            cond2 = ((listatesti['mappa'] - listatesti[f'mappa_{cod}']) != 0) & ((listatesti['mappa'] - listatesti[f'mappa_{cod}']).abs() <= intervallo) if ordinate == 0 else ((listatesti[f'mappa_{cod}'] - listatesti['mappa']) > 0) & ((listatesti[f'mappa_{cod}'] - listatesti['mappa']).abs() <= intervallo)
            cod += 1
            listatesti = listatesti[cond1 & cond2]

            if listatesti.empty:
                return []
            else:
                contexts = self.findcontexts(listatesti)
                bibliocontexts = self.findbib(contexts)
                clean = bibliocontexts.drop_duplicates(subset="contesto")
                highlights = self.highlight(clean)

        return highlights.to_dict(orient='records')
    
    ###################
    #   OLD
    ###################
    '''def ricerca_cooccorrenze (self, listaricerche, intervallo, periodo, ordinate):
        listatesti = pd.DataFrame()
        cod = 1
        if listaricerche[0][1] == 0:
            ricerca = self.sendBasicQuery(listaricerche[0][0], 'forma', listaricerche[0][2], listaricerche[0][3], pandas=True)
            listatesti = self.findtexts(0, ricerca)
        elif listaricerche[0][1] == 1:
            ricerca = self.sendBasicQuery(listaricerche[0][0], 'lemma', listaricerche[0][2], listaricerche[0][3], pandas=True)
            listatesti = self.findtexts(1, ricerca)
        elif listaricerche[0][1] == 2:
            ricerca = self.sendBasicQuery(listaricerche[0][0], 'lemma', listaricerche[0][2], listaricerche[0][3], pandas=True)
            listatesti = self.findtexts(2, ricerca)
        
        if listatesti.empty:
            return []

        for ricerca, tipo, espansa, raddoppiata in listaricerche[1:]:
            if tipo == 0:
                search = self.sendBasicQuery(ricerca, 'forma', espansa, raddoppiata, pandas=True)
            elif tipo == 1:
                search = self.sendBasicQuery(ricerca, 'lemma', espansa, raddoppiata, pandas=True)
            elif tipo == 2:
                search = self.sendBasicQuery(ricerca, 'lemma', espansa, raddoppiata, pandas=True)
     
            textlist = self.findtexts(tipo, search)
            cod_cols = []

            ##########################################
            # KORA: questo blocco sembra troppo lento!
            ##########################################
#             df_new = pd.DataFrame(columns=list(listatesti.columns))
#             for index1, row1 in listatesti.iterrows():
#                 for index2, row2 in textlist.iterrows():
#                     cond1 = row1['ntx'] == row2['ntx']
#                     cond2 = row1['numperiod'] == row2['numperiod'] if periodo == 1 else True
#                     cond3 = ((row1['mappa'] - row2['mappa']) != 0) and ((row1['mappa'] - row2['mappa']) in range(-intervallo, intervallo)) if ordinate == 0 else ((row2['mappa'] - row1['mappa']) > 0) and ((row2['mappa'] - row1['mappa']) <= intervallo)

#                     if cond1 and cond2 and cond3:
#                         row1[f'cod{cod}'] = textlist['cod'].iloc[index2]
#                         cod_cols.append(f'cod{cod}')
#                         df_new = pd.concat([df_new, row1.to_frame().T])

            ##########################################
            # KORA: fino a qui
            ##########################################

            ##########################################
            # KORA: questa è un'alternativa
            ##########################################
            df_new = pd.DataFrame(columns=list(listatesti.columns))
            for index1, row1 in listatesti.iterrows():
                ntx1 = row1['ntx']
                mappa1 = row1['mappa']
                nperiodo1 = row1['numperiod']
                if ordinate == 0: 
                    test = (textlist['ntx'].astype(int) == ntx1) & (textlist['mappa'].astype(int) > mappa1-intervallo) & (textlist['mappa'].astype(int) < mappa1+intervallo) & (textlist['mappa'].astype(int) != mappa1)
                else:
                    test = (textlist['ntx'].astype(int) == ntx1) & (textlist['mappa'].astype(int) > mappa1-intervallo)
                if periodo == 1:
                    test = test & (textlist['numperiod'] == nperiodo1)
                df_temp = textlist[test]
                df_temp[f'cod{cod}'] = df_temp['cod']
                cod_cols.append(f'cod{cod}')
                df_new = pd.concat([df_new, df_temp])
                cod += 1


            ##########################################
            # KORA: fino a qui
            ##########################################


            listatesti = df_new


            if listatesti.empty:
                return []
            else:
                contexts = self.findcontexts(listatesti)
                bibliocontexts = self.findbib(contexts)
                clean = bibliocontexts.drop_duplicates(subset="contesto")
                highlights = self.highlight (clean)
        return  highlights.to_dict(orient='records') # is this slow? CHECK!
        #return clean.to_json(orient='records') # possibilità alternativa -- molte opzioni possibili!'''