from typing import List, Tuple
from typing_extensions import Literal
import logging
import pandas as pd
from pandas import DataFrame, Series
from utils.config import getconfig
from utils.preprocessing import processingpipeline
import streamlit as st
from transformers import pipeline

@st.cache_resource
def load_adapmitClassifier(config_file:str = None, classifier_name:str = None):
    """
    loads the document classifier using haystack, where the name/path of model
    in HF-hub as string is used to fetch the model object.Either configfile or 
    model should be passed.
    1. https://docs.haystack.deepset.ai/reference/document-classifier-api
    2. https://docs.haystack.deepset.ai/docs/document_classifier
    Params
    --------
    config_file: config file path from which to read the model name
    classifier_name: if modelname is passed, it takes a priority if not \
    found then will look for configfile, else raise error.
    Return: document classifier model
    """
    if not classifier_name:
        if not config_file:
            logging.warning("Pass either model name or config file")
            return
        else:
            config = getconfig(config_file)
            classifier_name = config.get('adapmit','MODEL')
    
    logging.info("Loading Adaptation Mitigation classifier")    
    doc_classifier = pipeline("text-classification", 
                            model=classifier_name, 
                            return_all_scores=True, 
                            function_to_apply= "sigmoid")
    return doc_classifier


@st.cache_data
def adapmit_classification(haystack_doc:pd.DataFrame,
                        threshold:float = 0.5, 
                        classifier_model:pipeline= None
                        )->Tuple[DataFrame,Series]:
    """
    Text-Classification on the list of texts provided. Classifier provides the 
    most appropriate label for each text. these labels are in terms of if text 
    belongs to which particular Sustainable Devleopment Goal (SDG).
    Params
    ---------
    haystack_doc: List of haystack Documents. The output of Preprocessing Pipeline 
    contains the list of paragraphs in different format,here the list of 
    Haystack Documents is used.
    threshold: threshold value for the model to keep the results from classifier
    classifiermodel: you can pass the classifier model directly,which takes priority
    however if not then looks for model in streamlit session.
    In case of streamlit avoid passing the model directly.
    Returns
    ----------
    df: Dataframe
    """
    logging.info("Working on Adaptation-Mitigation Identification")
    haystack_doc['Adapt-Mitig Label'] = 'NA'
    haystack_doc['cond_check'] = haystack_doc.apply(lambda x: True if (
                (x['Target Label'] == 'TARGET') | (x['Action Label'] == 'Action') |
                (x['Policies_Plans Label'] == 'Policies and Plans')) else
                False, axis=1)
    # we apply Netzero to only paragraphs which are classified as 'Target' related
    df1 = haystack_doc[haystack_doc['cond_check'] == True]
    df1 = df1.reset_index(drop=True)
    df = haystack_doc[haystack_doc['cond_check'] == False]
    df = df.reset_index(drop=True)

    if not classifier_model:
        classifier_model = st.session_state['adapmit_classifier']
    
    predictions = classifier_model(list(df1.text))
     # converting the predictions to desired format
    list_ = []
    for i in range(len(predictions)):

      temp = predictions[i]
      placeholder = {}
      for j in range(len(temp)):
        placeholder[temp[j]['label']] = temp[j]['score']
      list_.append(placeholder)
    labels_ = [{**list_[l]} for l in range(len(predictions))]
    truth_df = DataFrame.from_dict(labels_)
    truth_df = truth_df.round(2)
    # convert the labels score into boolean based on threshold value
    truth_df = truth_df.astype(float) >= threshold
    truth_df = truth_df.astype(str)
    # list of labels
    categories = list(truth_df.columns)

    # collecting the labels, None is passed to overcome comprehension syntax
    truth_df['Adapt-Mitig Label'] = truth_df.apply(lambda x: {i if x[i]=='True' 
                                        else None for i in categories}, axis=1)
    truth_df['Adapt-Mitig Label'] = truth_df.apply(lambda x: 
                                  list(x['Adapt-Mitig Label'] -{None}),axis=1)
    # adding Adaptation-Mitigation label                              
    df1['Adapt-Mitig Label'] = list(truth_df['Adapt-Mitig Label'])
    df = pd.concat([df,df1])
    df = df.drop(columns = ['cond_check'])
    df = df.reset_index(drop =True)
    df.index += 1

    return df