# %%
import pandas as pd

from .basic_queries import basicQueries
from .utilities.format import formatAllContexts


# 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 di arrays del tipo:
    # [forma/lemma_cercati, tipo_ricerca, ricerca_espansa, iniziale_raddoppiata].
    # l'attributo tipo_ricerca ha come valori ammessi: 0 per forme, 1 per lemmi, 2 per lemmi + 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):
        occurrences = [] # una lista di Dicts con i codici dei lemmi/forme da cercare
        for ricerca, tipo, espansa, raddoppiata in listaricerche:
            if tipo==0:
                res1 = self.sendBasicQuery(ricerca, 'forma', espansa, raddoppiata, pandas=True)
                if res1.empty:
                    return []
                occurrences.append({'codList': list(res1['cod']), 'querySubtype': 0})
            elif tipo==1:
                res1 = self.sendBasicQuery(ricerca, 'lemma', espansa, raddoppiata, pandas=True)
                if res1.empty:
                    return []
                occurrences.append({'codList': list(res1['cod']), 'querySubtype': 1})
            elif tipo==2:
                res1 = self.sendBasicQuery(ricerca, 'lemma', espansa, raddoppiata, pandas=True)
                if res1.empty:
                    return []
                codList = list(res1['cod'])
                subQueryData = {'queryType': 'pfl', 'codList': codList}
                subdf = self.queryHandler.query(subQueryData, pandas=True)
                formCodList = list(subdf['codForma'])
                occurrences.append({'codList': codList, 'formCodList': formCodList, 'querySubtype': 2})

        if len(occurrences)==0:
            return []

        queryData = {'queryType': 'co-occurrences', 'occurrences': occurrences, 'intervallo': intervallo, 'periodo': periodo, 'ordinate': ordinate}
        queryResponses = [self.queryHandler.query(dict(queryData, table=table), pandas=True) for table in self.listOcc]
        listatesti = pd.concat(queryResponses)

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

        return highlights.to_dict(orient='records')