from typing import Literal
import a2.utils.utils
import a2.training.training_hugging
import transformers


SUPPORTED_MODELS: Literal = ["deberta_base", "deberta_small", "electra_base"]
SUPPORTED_TRAINERS: Literal = ["default", None]
SUPPORTED_LOSSES: Literal = ["default", None, "focal_loss"]

def get_model_config(model_name: SUPPORTED_MODELS):
    if model_name == "deberta_base" or model_name == "deberta_small":
        hyper_parameters = a2.training.training_hugging.HyperParametersDebertaClassifier()
    elif model_name == "electra_base":
        hyper_parameters = a2.training.training_hugging.HyperParametersElectraClassifier()
    else:
        raise ValueError(f"{model_name=} not supported, ({SUPPORTED_MODELS=})!")
    return hyper_parameters


def get_customized_trainer_class(trainer_name: SUPPORTED_TRAINERS, method_overwrites: list | None = None):
    if method_overwrites is None:
        method_overwrites = [] 
    if trainer_name == "default" or trainer_name is None:
        trainer = transformers.Trainer
    else:
        raise ValueError(f"{trainer_name=} not supported ({SUPPORTED_MODELS=})!")
    
    if "focal_loss" in method_overwrites:
        method_overwrites.remove("focal_loss")
        trainer.compute_loss = a2.training.training_hugging.TrainerWithFocalLoss.compute_loss

    if len(method_overwrites) != 0:
        raise ValueError(f"{method_overwrites=} not supported ({SUPPORTED_LOSSES=})!")

    return trainer