from abc import ABC
from typing import Optional
from transformers.trainer_utils import TrainOutput
from transformers.models.bert.modeling_bert import BertForSequenceClassification
from datasets import ClassLabel # type: ignore
import pandas as pd

from artifex.core import auto_validate_methods, ClassificationClassName, ValidationError
from artifex.models.classification_model import ClassificationModel
from artifex.config import config


@auto_validate_methods
class NClassClassificationModel(ClassificationModel, ABC):
    """
    A classification model in which the number of possible labels is not known upfront.
    """
    
    def __init__(self):
        super().__init__()
        # Labels are initialized to an empty ClassLabel, as the number of classes is not known upfront.
        self._labels_val: ClassLabel = ClassLabel(names=[])
        # Model is initialized to None, as the number of classes is not known upfront.
        self._model_val: Optional[BertForSequenceClassification] = None
        
    ##### Properties #####
    
    @property
    def _labels(self) -> ClassLabel:
        return self._labels_val
    
    @_labels.setter
    def _labels(self, labels: ClassLabel) -> None:
        self._labels_val = labels
    
    ##### Methods #####
    
    def _cleanup_synthetic_dataset(self, synthetic_dataset_path: str) -> None:
        """
        Remove from the synthetic training dataset:
        - All rows whose last element (the label) is not one of the accepted labels (the ones in self._labels).
        - All rows whose first element (the text) is shorter than 10 characters or is empty.
        """
        
        df = pd.read_csv(synthetic_dataset_path) # type: ignore
        valid_labels = set(self._labels.names)
        df = df[df.iloc[:, -1].isin(valid_labels)] # type: ignore
        df = df[df.iloc[:, 0].str.strip().str.len() >= 10] # type: ignore
        df.to_csv(synthetic_dataset_path, index=False)
    
    def train(
        self, classes: dict[str, str], output_path: Optional[str] = None, 
        num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3
    ) -> TrainOutput:
        f"""
        Train the IntentClassifier model using synthetic data generated by Synthex.
        Args:
            classes (dict[str, str]): A dictionary mapping class names to their descriptions. The keys 
                (class names) must be string with no spaces and a maximum length of 
                {config.INTENT_CLASSIFIER_CLASSNAME_MAX_LENGTH} characters.
            output_path (Optional[str]): The path where the generated synthetic data will be saved.
            num_samples (int): The number of training data samples to generate.
            num_epochs (int): The number of epochs for training the model.
        """
        
        # Validate class names, raise a ValidationError if any class name is invalid
        validated_classes: dict[ClassificationClassName, str] = {}
        for class_name, description in classes.items():
            try:
                validated_class_name = ClassificationClassName(class_name)
                validated_classes[validated_class_name] = description
            except ValueError:
                raise ValidationError(
                    message=f"`classes` keys must be strings with no spaces and a maximum length of {config.INTENT_CLASSIFIER_CLASSNAME_MAX_LENGTH} characters.",
                )

        # Populate the labels property with the validated class names
        validated_classnames = validated_classes.keys()
        self._labels = ClassLabel(names=list(validated_classnames))
        # Create the model with the correct number of labels
        self._model = BertForSequenceClassification.from_pretrained( # type: ignore
            config.INTENT_CLASSIFIER_HF_BASE_MODEL, num_labels=len(validated_classnames)
        )

        # Turn the validated classes into a list of instructions
        user_instructions: list[str] = self._parse_user_instructions(validated_classes)
        
        output: TrainOutput = self._train_pipeline(
            user_instructions=user_instructions, output_path=output_path, num_samples=num_samples, 
            num_epochs=num_epochs
        )
        
        return output