"""
Functions for performing inference and prediction using transformer models
"""
from transformers import AutoModelForSequenceClassification, AutoModelForTokenClassification, Trainer, pipeline
from datasets import Dataset
import bilbystats as bs
import numpy as np
import pandas as pd
from transformers import AutoTokenizer, DataCollatorWithPadding, DataCollatorForTokenClassification
from importlib import resources
from dotenv import load_dotenv
import os
import bilbystats as bs

# Use importlib to get a temporary path to the installed .env file
env_path = resources.files(
    "bilbystats.defaults").joinpath("local_defaults.env")

# Load environment variables from the .env file
load_dotenv(dotenv_path=env_path)

def load_model(model_path: str, model_name: str, model_type = 'classification', pad_to_multiple_of = 8):
    """
    Load a pretrained sequence classification model, its tokenizer, and initialize a Trainer.

    This function loads a Hugging Face `AutoModelForSequenceClassification` model from the specified 
    path, along with its corresponding tokenizer. It also initializes a `Trainer` instance using the model. 
    A temporary directory named 'tmp_trainer' is removed as part of the loading process.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    model_path : str
        Path to the pretrained model directory.
    model_name : str
        Identifier for the tokenizer to be loaded, typically matching a pretrained model name.

    ---------------------------------------------------------------------------
    OUTPUT:
    model_dict : dict
        A dictionary containing the loaded components:
            - 'model': AutoModelForSequenceClassification
            - 'trainer': Trainer
            - 'tokenizer': AutoTokenizer
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """ 
    try: 
        tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
    except:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        
    if model_type == 'classification':
        model = AutoModelForSequenceClassification.from_pretrained(model_path)
        collator  = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=pad_to_multiple_of)
        trainer = Trainer(model=model, data_collator=collator)
    elif model_type == 'entities':
        model = AutoModelForTokenClassification.from_pretrained(model_path)
        collator  = DataCollatorForTokenClassification(tokenizer, pad_to_multiple_of=pad_to_multiple_of)
        trainer = Trainer(model=model, data_collator=collator)
        
    model_dict = {
        "model": model,
        "trainer": trainer,
        "tokenizer": tokenizer,
        "model_type": model_type
        }
    
    os.system("rm -rf tmp_trainer")
    
    return model_dict

def predict(data: Dataset, model_dict, padding = False, tokenize = True) -> dict:
    """
    Generate predictions using a fine-tuned transformer model on input text data.

    This function handles string, list, or Dataset input types by converting them into 
    a Hugging Face Dataset format. It tokenizes the data using a specified tokenizer, 
    loads a sequence classification model from the given path, and runs inference using 
    Hugging Face's `Trainer`. The output includes raw logits, predicted labels, and 
    optionally true labels if available.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    data : Dataset, list, or str
        Input data to predict on. May be:
            - a Hugging Face Dataset with a 'text' field,
            - a string (single text input), or
            - a list of strings.
        If true labels are included, they should be in a 'label' column.

    model_path : str
        Path to the directory containing the fine-tuned transformer model.

    model_name : str
        Name or path of the model used to initialize the tokenizer.

    ---------------------------------------------------------------------------
    OUTPUT:
    output : dict
        A dictionary containing:
            - 'logits': np.ndarray
                The raw output scores from the model.
            - 'pred_labels': np.ndarray
                The predicted class labels based on the logits.
            - 'true_labels': np.ndarray, optional
                Ground truth labels if present in the input data.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    trainer = model_dict['trainer']
    tokenizer = model_dict['tokenizer']

    # If the datatype is a string convert to a dataset for evaluation
    if type(data) == str:
        data = Dataset.from_dict({'text': [data]})
    elif type(data) == list:
        data = Dataset.from_dict({'text': data})

    # Tokenize the input data
    if tokenize:
        data2predict_tk = bs.tokenize(data, tokenizer, padding = padding)
    else:
        data2predict_tk = data
        
    # Calculate the model predictions
    predictions = trainer.predict(data2predict_tk)

    model_type = model_dict['model_type']
    if model_type == 'classification':
        axis = 1
    elif model_type == 'entities':
        axis = 2

    pred_labels = np.argmax(predictions.predictions, axis=axis)

    # Save the output as a dictionary
    output = {
        'logits': predictions.predictions,
        'pred_labels': pred_labels
    }

    if hasattr(predictions, 'label_ids') and predictions.label_ids is not None:
        output['true_labels'] = predictions.label_ids

    return output


def predict_df(df: pd.DataFrame, covariate: str, target: str, model_dict,
                                indices: list[int] = None) -> dict:
    """
    Generate predictions from a model for a DataFrame subset based on specified covariate.

    This function selects a subset of the input DataFrame (if indices are provided),
    converts the data into a format suitable for model inference, and uses a fine-tuned
    transformer model to compute predictions. Optionally includes the true target labels.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    df : pd.DataFrame
        Input DataFrame containing the covariate column used for prediction.
    covariate : str
        Name of the column in `df` containing the input text data.
    model_path : str
        Path to the directory containing the fine-tuned transformer model.
    model_name : str
        Name or path of the pre-trained model for tokenizer loading.
    indices : list of int, optional
        Indices specifying the subset of rows from `df` to predict on.
        If None, the entire DataFrame is used.
    target : str, optional
        Name of the column in `df` containing target labels, if available.

    ---------------------------------------------------------------------------
    OUTPUT:
    output : dict
        A dictionary containing:
            - 'logits': np.ndarray, raw model output scores.
            - 'pred_labels': np.ndarray, predicted class labels.
            - 'true_labels': np.ndarray (optional), ground-truth labels if provided.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if indices is not None:
        subset_df = df.iloc[indices].copy()
    else:
        subset_df = df.copy()

    data2predict = bs.df2hfdict(subset_df, covariate, target)

    output = bs.predict(data2predict, model_dict)

    return output

def load_pretrained_model(my_model_name):
    if my_model_name == 'sentence_sentiment':
        model_name = "hfl/chinese-roberta-wwm-ext"
        model_path = os.getenv("SENTENCE_SENTIMENT_CHECKPOINT")
        model_dict = load_model(model_path, model_name)

    return model_dict

def ner_pipeline(model_dict, aggregation_strategy = 'first', max_length = 512):
    tokenizer=model_dict['tokenizer']

    if max_length:
        tokenizer.model_max_length = max_length
    
    out = pipeline("ner", model=model_dict['model'], tokenizer = tokenizer, aggregation_strategy=aggregation_strategy)
    
    return out

def docs2entities(documents, model_dict, ner_batch_size = 16, progress_size = 50, aggregation_strategy = 'first', doc_ids = None, language = 'English'):
    nonnonedocs = [x for x in documents if x is not None]
    if not doc_ids:
        nonnoneids = [None] * len(nonnonedocs)
    else:
        nonnoneids = bs.filter_nones(documents, doc_ids)
        
    #document_sentences_perdoc_orig, _ = bs.docs2sentences(nonnonedocs, language = 'English', combine = False)
    if language == 'English':
        nonnonedocs = bs.rep_texts(nonnonedocs, '\n', ' ')
        
    document_sentences, doclengths = bs.docs2sentences(nonnonedocs, language = language, combine = True)
    document_sentences_perdoc, _ = bs.docs2sentences(nonnonedocs, language = language, combine = False)
    ner_pipe = bs.ner_pipeline(model_dict, aggregation_strategy)
    
    try:
        environment_stage = os.environ["environment"]
        if environment_stage.lower() == 'production' or  environment_stage.lower() == 'staging':
            progress_size = 0 
    except:
        pass
    
    if progress_size == 0:
        entity_set = ner_pipe(document_sentences, batch_size=ner_batch_size)
    else:
        entity_set = []
        total_sentences = len(document_sentences)
        for i in range(0, total_sentences, progress_size):
            bs.loader(i, total_sentences, desc="Processing NER")
            batch_end = min(i + progress_size, total_sentences)
            batch_sentences = document_sentences[i:batch_end]
            entity_set_new = ner_pipe(batch_sentences, batch_size=ner_batch_size)
            entity_set.extend(entity_set_new)
    
    document_entities = bs.parse_list(entity_set, doclengths)
    document_entities_processed = []
    for i in range(len(nonnonedocs)):
        doc_id = nonnoneids[i]
        output = bs.document_entities(document_entities[i], document_sentences_perdoc[i], doc_id = doc_id)
        #output = bs.document_entities(document_entities[i], document_sentences_perdoc_orig[i], doc_id = doc_id)
        document_entities_processed.append(output)
            
    out = bs.none_list_comb(documents, document_entities_processed)
    
    return out

