"""Module for artifact and trial management in the pipeline.

This module defines data models for managing artifacts and trials in the pipeline,
including their configurations, outputs, and optimization details.
"""

from typing import Any

import numpy as np
from numpy.typing import NDArray
from pydantic import BaseModel, ConfigDict, Field

from autointent.configs import EmbedderConfig
from autointent.custom_types import ListOfLabelsWithOOS, NodeType


class Artifact(BaseModel):
    """Base class for artifacts generated by pipeline nodes."""


class RegexArtifact(Artifact):
    """Artifact containing results from the regex node."""


class EmbeddingArtifact(Artifact):
    """Artifact containing details from the embedding node.

    Attributes:
        config: Configuration settings for the embedder.
    """

    config: EmbedderConfig


class ScorerArtifact(Artifact):
    """Artifact containing outputs from the scoring node.

    Outputs from the best scorer, numpy arrays of shape (n_samples, n_classes).

    Attributes:
        train_scores: Scorer outputs for train utterances.
        validation_scores: Scorer outputs for validation utterances.
        test_scores: Scorer outputs for test utterances.
        folded_scores: Scores for each fold from cross-validation.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)
    train_scores: NDArray[np.float64] | None = Field(None, description="Scorer outputs for train utterances")
    validation_scores: NDArray[np.float64] | None = Field(None, description="Scorer outputs for validation utterances")
    test_scores: NDArray[np.float64] | None = Field(None, description="Scorer outputs for test utterances")
    folded_scores: list[NDArray[np.float64]] | None = Field(
        None, description="Scores for each fold from cross-validation"
    )


class DecisionArtifact(Artifact):
    """Artifact containing outputs from the predictor node.

    Outputs from the best predictor, numpy array of shape (n_samples,) or
    (n_samples, n_classes) depending on classification mode (multi-class or multi-label).

    Attributes:
        labels: Predicted labels for the samples.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)
    labels: ListOfLabelsWithOOS


def validate_node_name(value: str) -> str:
    """Validate and return the node type.

    Args:
        value: Node type as a string.

    Returns:
        Validated node type string.

    Raises:
        ValueError: If the node type is invalid.
    """
    if value in [NodeType.embedding, NodeType.scoring, NodeType.decision, NodeType.regex]:
        return value
    msg = f"Unknown node_type: {value}. Expected one of ['regex', 'embedding', 'scoring', 'decision']"
    raise ValueError(msg)


class Artifacts(BaseModel):
    """Container for storing and managing artifacts generated by pipeline nodes.

    Modules hyperparams and outputs. The best ones are transmitted between nodes of the pipeline.

    Attributes:
        regex: List of artifacts from the regex node.
        embedding: List of artifacts from the embedding node.
        scoring: List of artifacts from the scoring node.
        decision: List of artifacts from the decision node.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    regex: list[RegexArtifact] = []
    embedding: list[EmbeddingArtifact] = []
    scoring: list[ScorerArtifact] = []
    decision: list[DecisionArtifact] = []

    def add_artifact(self, node_type: str, artifact: Artifact) -> None:
        """Add an artifact to the specified node type.

        Args:
            node_type: Node type as a string.
            artifact: The artifact to add.
        """
        self.get_artifacts(node_type).append(artifact)

    def get_artifacts(self, node_type: str) -> list[Artifact]:
        """Retrieve all artifacts for a specified node type.

        Args:
            node_type: Node type as a string.

        Returns:
            A list of artifacts for the node type.
        """
        return getattr(self, validate_node_name(node_type))  # type: ignore[no-any-return]

    def get_best_artifact(self, node_type: str, idx: int) -> Artifact:
        """Retrieve the best artifact for a specified node type and index.

        Args:
            node_type: Node type as a string.
            idx: Index of the artifact.

        Returns:
            The best artifact.
        """
        return self.get_artifacts(node_type)[idx]


class Trial(BaseModel):
    """Representation of an individual optimization trial.

    Attributes:
        module_name: Type of the module being optimized.
        module_params: Parameters of the module for the trial.
        metric_name: Name of the evaluation metric.
        metric_value: Value of the evaluation metric for this trial.
        module_dump_dir: Directory where the module is dumped.
        metrics: Dictionary of metric names and their values.
    """

    module_name: str
    module_params: dict[str, Any]
    metric_name: str
    metric_value: float
    module_dump_dir: str | None
    metrics: dict[str, float]


class Trials(BaseModel):
    """Container for managing optimization trials for pipeline nodes.

    Attributes:
        regex: List of trials for the regex node.
        embedding: List of trials for the embedding node.
        scoring: List of trials for the scoring node.
        decision: List of trials for the decision node.
    """

    regex: list[Trial] = []
    embedding: list[Trial] = []
    scoring: list[Trial] = []
    decision: list[Trial] = []

    def get_trial(self, node_type: str, idx: int) -> Trial:
        """Retrieve a specific trial for a node type and index.

        Args:
            node_type: Node type as a string.
            idx: Index of the trial.

        Returns:
            The requested trial.
        """
        return self.get_trials(node_type)[idx]

    def get_trials(self, node_type: str) -> list[Trial]:
        """Retrieve all trials for a specified node type.

        Args:
            node_type: Node type as a string.

        Returns:
            A list of trials for the node type.
        """
        return getattr(self, validate_node_name(node_type))  # type: ignore[no-any-return]

    def add_trial(self, node_type: str, trial: Trial) -> None:
        """Add a trial to a specified node type.

        Args:
            node_type: Node type as a string.
            trial: The trial to add.
        """
        self.get_trials(node_type).append(trial)


class TrialsIds(BaseModel):
    """Representation of the best trial IDs for each pipeline node.

    Attributes:
        regex: Best trial index for the regex node.
        embedding: Best trial index for the embedding node.
        scoring: Best trial index for the scoring node.
        decision: Best trial index for the decision node.
    """

    regex: int | None = None
    embedding: int | None = None
    scoring: int | None = None
    decision: int | None = None

    def get_best_trial_idx(self, node_type: str) -> int | None:
        """Retrieve the best trial index for a specified node type.

        Args:
            node_type: Node type as a string.

        Returns:
            The index of the best trial, or None if not set.
        """
        return getattr(self, validate_node_name(node_type))  # type: ignore[no-any-return]

    def set_best_trial_idx(self, node_type: str, idx: int) -> None:
        """Set the best trial index for a specified node type.

        Args:
            node_type: Node type as a string.
            idx: Index of the best trial.
        """
        setattr(self, validate_node_name(node_type), idx)
