#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed May 14 11:43:28 2025

@author: samd
"""

import torch
import os
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification, Trainer, TrainingArguments, DataCollatorWithPadding, DataCollatorForTokenClassification

import logging  # Python's built-in logging
import transformers.utils.logging as transformers_logging  # Transformers logging

from tqdm import tqdm
import os
import bilbystats as bs
from datasets import Dataset
import numpy as np
import warnings
import wandb
from datetime import datetime

from importlib import resources
from dotenv import load_dotenv
import os

import subprocess
import sys

import torch

class MultiLabelDataCollator:
    def __call__(self, features):
        batch = {
            'input_ids': torch.tensor([f['input_ids'] for f in features], dtype=torch.long),
            'attention_mask': torch.tensor([f['attention_mask'] for f in features], dtype=torch.long),
            'labels': torch.tensor([f['labels'] for f in features], dtype=torch.float)
        }
        return batch

def trainTFmodel(
    train_data: Dataset,
    valid_data: Dataset,
    model_name: str,
    savename: str = None,
    savedir: str = "./",
    num_labels: int = 2,
    label2id: dict = None,
    training_args: TrainingArguments = None,
    use_wandb: bool = True,
    eval_strat: str = "steps",
    nepochs = None,
    padding = False,
    model_type = None
) -> tuple[Trainer, AutoModelForSequenceClassification, TrainingArguments]:
    """
    Train a transformer model for sequence classification using Hugging Face Trainer.

    This function initializes and trains a transformer-based model using the Hugging Face 
    `Trainer` API. It supports optional custom label mappings and training arguments, 
    and disables Weights & Biases logging by default.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    train_data : Dataset
        Tokenized training dataset compatible with Hugging Face models.
    valid_data : Dataset
        Tokenized validation dataset for evaluation during training.
    model_name : str
        Name or path of the pre-trained model to fine-tune.
    savename : str, optional
        Optional name for saving the trained model outputs.
    savedir : str, optional (default="./")
        Directory path where model outputs will be saved.
    num_labels : int, optional (default=2)
        Number of unique class labels for classification.
    label2id : dict, optional
        Dictionary mapping label names to integer IDs.
    training_args : TrainingArguments, optional
        Hugging Face TrainingArguments object. If not provided, default settings are used.
    use_wandb: Bool, optional
        An indictor of whether to use wandb for tracking. Default is True.
    ---------------------------------------------------------------------------
    OUTPUT:
    trainer : Trainer
        The Hugging Face Trainer instance used for model training.
    model : AutoModelForSequenceClassification
        The trained transformer classification model.
    training_args : TrainingArguments
        The training arguments used during training.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if not os.path.isdir(savedir):
        raise ValueError('savedir is not a valid directory')
        return
    
    if use_wandb:
        os.environ["TOKENIZERS_PARALLELISM"] = "false"
        wandb.init(project=savename,
               name=f"training_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
               settings=wandb.Settings(quiet=True))

    logging.getLogger("transformers.modeling_utils").setLevel(logging.ERROR)
    
    if training_args == 'entities' and not model_type:
        model_type = 'entities'
        
    if not label2id:
        label2id = make_mapping()
        
    id2label = bs.invert_dict(label2id)
    if not model_type:
        model_type = 'classification'
    
    tokenizer = bs.load_tokenizer(model_name)

    if model_type == 'classification':
        compute_metric_fn = bs.compute_metrics
    
        if not label2id:
            model, loading_info = AutoModelForSequenceClassification.from_pretrained(
                model_name, num_labels=num_labels, output_loading_info=True)
        else:
            if len(label2id) != num_labels:
                raise ValueError("The number of label ids does not match the number of labels")

            # Load the model
            model, loading_info = AutoModelForSequenceClassification.from_pretrained(
                model_name,
                num_labels=num_labels,
                id2label=id2label,
                label2id=label2id,
                output_loading_info=True)
        
            
        data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
    
    elif model_type == 'multi_label_classification':
        compute_metric_fn = bs.compute_metrics_multilabel
    
        model, loading_info = AutoModelForSequenceClassification.from_pretrained(
            model_name, num_labels=num_labels, output_loading_info=True, problem_type="multi_label_classification")
            
        data_collator = MultiLabelDataCollator()
            
    else:
        compute_metric_fn = bs.make_entity_metrics(id2label)
        
        model = AutoModelForTokenClassification.from_pretrained(
                model_name,
                num_labels=len(label2id),
                id2label=id2label,
                label2id=label2id)
        
        data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer, pad_to_multiple_of=8)

    # Set WANDB API key
    if use_wandb:
        try:
            os.environ["WANDB_API_KEY"] = bs.read_api_key('WANDB')
        except:
            if "WANDB_API_KEY" in os.environ:
                warnings.warn("WANDB_API_KEY not set in the package API keys but found in the exported environment. Now proceeding using the discovered key.")
            else:
                warnings.warn("WANDB_API_KEY not set, please export it or add it to the API_keys file as instructed in the README and reinstall the package. For now reverting to usage without wandb.")
                use_wandb = False
    
    # Set WANDB variables
    if use_wandb:
        os.environ["WANDB_PROJECT"] = savename
        os.environ["WANDB_LOG_MODEL"] = "false"
    else:
        os.environ["WANDB_DISABLED"] = "true"
        os.environ["WANDB_MODE"] = "offline"

    # Obtain the number of training samples
    ndata = len(train_data)
    
    # Obtain the training arguments
    if not training_args:
        training_args = default_training_args(model_name, savename, savedir, ndata = ndata, use_wandb = use_wandb, eval_strat = eval_strat, nepochs = nepochs)
    elif training_args == 'entities':
        training_args = default_training_args(model_name, savename, savedir, argfile = 'entities', ndata = ndata, use_wandb = use_wandb, eval_strat = eval_strat, nepochs = nepochs)
    else:
        training_args = training_args
        
        
    # Load the tokenizer to allow for dynamic padding
    if not padding:
        # Initialize the trainer
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_data,
            eval_dataset=valid_data,
            compute_metrics=compute_metric_fn,
            data_collator=data_collator,
        )
    else:
        # Initialize the trainer
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=train_data,
            eval_dataset=valid_data,
            compute_metrics=compute_metric_fn,
        )
        
    # Train the model
    trainer.train()
    
    # Exit from wandb without output
    wandb.finish()
    
    # Remove the saved wandb files
    os.system("rm -rf wandb/")

    return trainer, model, training_args


def default_training_args(model_name: str, savename: str = None,
                          savedir: str = './', argfile: str = None, ndata: int = 0, 
                          use_wandb: bool = True, eval_strat: str = "steps", nepochs = None) -> TrainingArguments:
    """
    Generate default training arguments for a Hugging Face model.

    This function sets up a default configuration of `TrainingArguments` used 
    for model training. If an `argfile` is not specified, it loads training 
    defaults from a predefined `.env` file. The model's output directory is 
    configured using the provided `savename` or falls back to `model_name`.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    model_name : str
        Name of the model to be used in the training configuration and run name.
    savename : str, optional
        Custom name for the training run. Defaults to `model_name` if not provided.
    savedir : str, optional
        Directory in which to save the training outputs. Defaults to current directory './'.
    argfile : str, optional
        Path to a file containing training arguments (currently unused but 
        triggers loading defaults from environment if not specified).
    ndata: int, optional
        The number of datapoints in the dataset
    ---------------------------------------------------------------------------
    OUTPUT:
    training_args : TrainingArguments
        An instance of Hugging Face's `TrainingArguments` initialized with default settings.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if not savename:
        savename = model_name

    if not argfile:
        # Use importlib to get a temporary path to the installed .env file
        env_path = resources.files(
            "bilbystats.defaults.training_defaults").joinpath("training_defaults.env")

        # Load environment variables from the .env file
        load_dotenv(dotenv_path=env_path)
    elif argfile == 'entities':
        env_path = resources.files(
            "bilbystats.defaults.training_defaults").joinpath("entity_training_defaults.env")
        load_dotenv(dotenv_path=env_path)
    else:
        load_dotenv(dotenv_path=argfile)

    def str_to_bool(value):
        return value.lower() in ("true", "1", "yes", "on")
    
    if ndata != 0:
        batch_size = int(
            os.getenv('PER_DEVICE_EVAL_BATCH_SIZE'))
        nloggingsteps = int(ndata/batch_size/10)
        nevalsteps = int(ndata/batch_size/2)
    else:
        nloggingsteps = 300
        nevalsteps = 300
        
    if use_wandb:
        rep2 = "wandb"
    else:
        rep2 = "none"
        
    if not nepochs:
        nepochs = int(os.getenv('NUM_TRAIN_EPOCHS'))

    cuda_available = torch.cuda.is_available()

    training_args = TrainingArguments(
        output_dir=savedir + savename,
        run_name=f"{model_name}",
        learning_rate=float(os.getenv('LEARNING_RATE')),
        do_eval=True,
        #eval_strategy=os.getenv('EVAL_STRATEGY'),
        eval_strategy = eval_strat,
        save_strategy = eval_strat,
        #save_strategy=os.getenv('SAVE_STRATEGY'),
        save_total_limit=int(os.getenv('SAVE_TOTAL_LIMIT')),
        auto_find_batch_size=str_to_bool(os.getenv('AUTO_FIND_BATCH_SIZE')),
        num_train_epochs=nepochs,
        weight_decay=float(os.getenv('WEIGHT_DECAY')),
        warmup_ratio=float(os.getenv('WARMUP_RATIO')),
        metric_for_best_model=str(os.getenv('METRIC_FOR_BEST_MODEL')),
        report_to=rep2,
        # how often to log to W&B
        eval_steps=nevalsteps,
        save_steps=nevalsteps,
        logging_steps=nloggingsteps,
        #logging_steps=int(os.getenv('LOGGING_STEPS')),
        load_best_model_at_end=True,
        seed=int(os.getenv('SEED')),
        data_seed=int(os.getenv('DATA_SEED')),
        save_safetensors=str_to_bool(os.getenv('SAVE_SAFETENSORS')),
        per_device_eval_batch_size=int(
            os.getenv('PER_DEVICE_EVAL_BATCH_SIZE')),
        per_device_train_batch_size=int(
            os.getenv('PER_DEVICE_TRAIN_BATCH_SIZE')),
        dataloader_pin_memory=cuda_available
    )
    return training_args


def default_training_args_dep(model_name: str, savename: str = None,
                              savedir: str = './') -> TrainingArguments:
    """
    Generate default training arguments for Hugging Face Trainer.

    This function creates a `TrainingArguments` object with sensible defaults 
    for fine-tuning a transformer model using the Hugging Face `Trainer` API. 
    It supports optional customization of output directory naming.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    model_name : str
        Name or identifier of the model, used for naming the run.
    savename : str, optional
        Custom name to use for the saved model output directory. 
        Defaults to `model_name` if not specified.
    savedir : str, optional (default='./')
        Base directory where training outputs will be saved.

    ---------------------------------------------------------------------------
    OUTPUT:
    training_args : TrainingArguments
        A configured TrainingArguments object ready for use with Hugging Face Trainer.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    if not savename:
        savename = model_name

    training_args = TrainingArguments(
        output_dir=savedir + savename,
        run_name=f"{model_name}",
        learning_rate=2e-5,
        do_eval=True,
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=1,
        auto_find_batch_size=True,
        num_train_epochs=3,
        weight_decay=0.01,
        logging_steps=300,  # how often to log to W&B
        load_best_model_at_end=True,
        # save_total_limit=1,
        seed=42,
        data_seed=42,
        save_safetensors=False,
    )
    return training_args

def make_mapping(strings):
    """
    Given a list of strings, return a mapping (dict) that maps each
    unique string to an integer from 0 to n-1.
    """
    unique_strings = list(dict.fromkeys(strings))  # preserves order, removes duplicates
    return {s: i for i, s in enumerate(unique_strings)}
