import gradio as gr
from os import environ
from transformers import pipeline

api_key = environ.get("api_key")

def app(input, filter_strategy, relevence_threshold, k, length_management_strategy, token_overlap, results_merge_stategy, drop_last):
    
    classifier_pipeline = pipeline(model="gianma/testModel", tokenizer='gianma/testModel', use_auth_token=api_key)
    length_limit = classifier_pipeline.tokenizer.model_max_length
    kwargs = {'padding':True,'truncation':True,'max_length':length_limit}

    kwargs['top_k'] = classifier_pipeline.model.num_labels

    if length_management_strategy=='solo inizio':
        res = classifier_pipeline(input, **kwargs)
        res = parse_results(res) # convert to expected format for gradio interface    
    elif length_management_strategy=='intero documento':
        res = execute_pipeline_su_chunks(classifier_pipeline, input, length_limit,drop_last = drop_last, token_overlap=token_overlap, merge_strategy =results_merge_stategy ,pipeline_kwargs=kwargs)
        

    if filter_strategy == 'soglia di confidenza':
        res = {k:v for k,v in res.items() if v >= relevence_threshold}
    elif filter_strategy == 'top k':
        rev_res = sorted([(v,k) for k,v in res.items()], reverse=True)
        res = {k:v for (v,k) in rev_res[:k] }

    return res


def execute_pipeline_su_chunks(pipeline, input_string, max_length, drop_last = False, token_overlap=0, merge_strategy ='avg' ,pipeline_kwargs={}):
    tokens = pipeline.tokenizer.tokenize(input_string)
    results = {}
    
    overlap_len = int(max_length * token_overlap)
    shift_len = max_length - overlap_len
    
    start = 0
    end = max_length

    iterations = 0

    while start < len(tokens):
        chunk_tokens = tokens[start:end]

        start += shift_len
        end = start + max_length

        is_last = not (start < len(tokens))
        
        if not (is_last and drop_last):
            chunk_string = pipeline.tokenizer.convert_tokens_to_string(chunk_tokens)
            chunck_res = pipeline(chunk_string, **pipeline_kwargs)
            chunck_res = parse_results(chunck_res) # convert to expected format for gradio interface   

            if results != {}:
                new_results = {}
                for k, v in results.items():
                    if merge_strategy == 'avg':
                        new_results[k] = v + chunck_res[k]
                    elif merge_strategy == 'max':
                        new_results[k] = max(v, chunck_res[k])
                results = new_results
            else:
                results = chunck_res

            iterations += 1
    
    if merge_strategy == 'avg':
        results = {k:v/iterations for k,v in results.items()}

    return results


def parse_results(res):
    return  {el['label']:el['score'] for el in res} # convert to expected format for gradio interface



with gr.Blocks() as interface:

    gr.Markdown("""
        # Classificatore domini EuroVoc

        Strumento che permette la classificazione di un documento rispetto ai 21 domini della classificazione EuroVoc. 
        """)

    with gr.Row():
      with gr.Column():

        input_text = gr.Textbox(label='Testo da classificare')

        with gr.Accordion("Opzioni"):

          with gr.Box():
            filter_strategy = gr.Radio(['soglia di confidenza', 'top k'], value='soglia di confidenza', visible=False)
            gr.Markdown("""
            ### Strategia di filtraggio dei risultati

            Indica se mostrare solamente i risultati sopra una certa soglia di confidenza oppure se mostrare le k classi più probabili
            """)
            with gr.Tab("soglia di confidenza") as tab_confidence:
              s_confidence = gr.Slider(0, 1, value=0.5,label='valore soglia',)
            with gr.Tab("top k") as tab_top_k:
              s_k = gr.Slider(1, 21, step= 1, value=4, label='k')
            tab_confidence.select(lambda :"soglia di confidenza", None, filter_strategy)
            tab_top_k.select(lambda:'top k', None, filter_strategy)

          with gr.Box():
            gr.Markdown("""
            ### Strategia di elaborazione del documento

            Indica se analizzare solamente la prima posrzione del documento o se eseguire più classificazioni rispetto alle diverse porzioni del documento unendone i risultati
            """)
            document_reading_strategy = gr.Radio(['solo inizio', 'intero documento'], value='solo inizio', visible=False)
            with gr.Tab("solo inizio") as tab_beginning_only:
              gr.Markdown("Opzione migliore come velocità di elaborazione")
            with gr.Tab("intero documento") as tab_whole_document:
              gr.Markdown("Attenzione: questa opzione allunga i tempi di elaborazione")
              chunk_overlap = gr.Slider(0, 0.5, value=0, label='sovrapposizione porzioni documento', info='indica la percentuale di sovrapposizione tra le diverse porzioni di testo analizzate')
              combine_strategy = gr.Radio(['avg', 'max'], value='avg', label='strategia per la combinazione dei risultati dei chunk', info='avg indica di considerare il valore medio delle classificazioni eseguite sulle diverse porzioni del documento, mentre max prenderà i valori massimi (aumenta il numero di classi rilevanti)')
              exclude_last = gr.Checkbox(value=True, label='escludi ultimo chuck', info='permette di rimuovere l\' ultima porzione del testo quando viene analizzato compretamente per evitare di elaborare un porzione troppo ridotta')
            tab_beginning_only.select(lambda:'solo inizio', None, document_reading_strategy)
            tab_whole_document.select(lambda:'intero documento', None, document_reading_strategy)

        with gr.Accordion('esempi'):
            gr.Examples(
                [["hi"], ["hello"]],
                [input_text],
            )
            
        with gr.Row():
          clear_button = gr.Button("Cancella testo")
          classify_button = gr.Button("Classifica", variant='primary')                        
        
      with gr.Column():
        outputs = gr.Label()

    clear_button.click(fn=lambda: None, inputs=None, outputs=input_text)
    classify_button.click(app, inputs=[input_text,filter_strategy, s_confidence, s_k, document_reading_strategy, chunk_overlap, combine_strategy, exclude_last], outputs=outputs)

interface.launch()