"""Module for managing pipeline optimization.

This module handles the tracking and logging of optimization artifacts,
trials, and modules during the pipeline's execution.
"""

import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

import numpy as np
from numpy.typing import NDArray

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

from ._data_models import Artifact, Artifacts, EmbeddingArtifact, ScorerArtifact, Trial, Trials, TrialsIds

if TYPE_CHECKING:
    from autointent.modules.base import BaseModule


@dataclass
class ModulesList:
    """Container for managing lists of modules for each node type.

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

    regex: list["BaseModule"] = field(default_factory=list)
    embedding: list["BaseModule"] = field(default_factory=list)
    scoring: list["BaseModule"] = field(default_factory=list)
    decision: list["BaseModule"] = field(default_factory=list)

    def get(self, node_type: str) -> list["BaseModule"]:
        """Retrieve the list of modules for a specific node type.

        Args:
            node_type: The type of node (e.g., "regex", "embedding").

        Returns:
            List of modules for the specified node type.
        """
        return getattr(self, node_type)  # type: ignore[no-any-return]

    def add_module(self, node_type: str, module: "BaseModule") -> None:
        """Add a module to the list for a specific node type.

        Args:
            node_type: The type of node.
            module: The module to add.
        """
        self.get(node_type).append(module)


class OptimizationInfo:
    """Tracks optimization results, including trials, artifacts, and modules.

    This class provides methods for logging optimization results, retrieving
    the best-performing modules and artifacts, and generating configuration
    for inference nodes.

    Attributes:
        artifacts: Container for storing optimization artifacts.
        trials: Container for storing optimization trials.
        modules: Container for storing module instances.
        pipeline_metrics: Dictionary storing pipeline-level metrics.
    """

    def __init__(self) -> None:
        """Initialize optimization info."""
        self._logger = logging.getLogger(__name__)

        self.artifacts = Artifacts()
        self.trials = Trials()
        self._trials_best_ids = TrialsIds()
        self.modules = ModulesList()
        self.pipeline_metrics: dict[str, float] = {}

    def log_module_optimization(
        self,
        node_type: str,
        module_name: str,
        module_params: dict[str, Any],
        metric_value: float,
        metric_name: str,
        metrics: dict[str, float],
        artifact: Artifact,
        module_dump_dir: str | None,
        module: "BaseModule | None" = None,
    ) -> None:
        """Log optimization results for a module.

        Args:
            node_type: Type of the node being optimized.
            module_name: Type of the module.
            module_params: Parameters of the module for the trial.
            metric_value: Metric value achieved by the module.
            metric_name: Name of the evaluation metric.
            metrics: Dictionary of metric names and their values.
            artifact: Artifact generated by the module.
            module_dump_dir: Directory where the module is dumped.
            module: The module instance, if available.
        """
        trial = Trial(
            module_name=module_name,
            metric_name=metric_name,
            metric_value=metric_value,
            module_params=module_params,
            module_dump_dir=module_dump_dir,
            metrics=metrics,
        )
        self.trials.add_trial(node_type, trial)
        self._logger.debug("module %s fitted and saved to optimization info", module_name, extra=trial.model_dump())

        if module:
            self.modules.add_module(node_type, module)

        self.artifacts.add_artifact(node_type, artifact)

    def _get_metrics_values(self, node_type: str) -> list[float]:
        """Retrieve all metric values for a specific node type.

        Args:
            node_type: Type of the node.

        Returns:
            List of metric values.
        """
        return [trial.metric_value for trial in self.trials.get_trials(node_type)]

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

        Args:
            node_type: Type of the node.

        Returns:
            Index of the best trial, or None if no trials exist.
        """
        if not self.trials.get_trials(node_type):
            return None
        best_idx = self._trials_best_ids.get_best_trial_idx(node_type)
        if best_idx is not None:
            return best_idx
        best_idx = int(np.argmax(self._get_metrics_values(node_type)))
        self._trials_best_ids.set_best_trial_idx(node_type, best_idx)
        return best_idx

    def _get_best_artifact(self, node_type: str) -> EmbeddingArtifact | ScorerArtifact | Artifact:
        """Retrieve the best artifact for a specific node type.

        Args:
            node_type: Type of the node.

        Returns:
            The best artifact for the node type.

        Raises:
            ValueError: If no best trial exists for the node type.
        """
        best_idx = self._get_best_trial_idx(node_type)
        if best_idx is None:
            msg = f"No best trial for {node_type}"
            raise ValueError(msg)
        return self.artifacts.get_best_artifact(node_type, best_idx)

    def get_best_embedder(self) -> EmbedderConfig:
        """Retrieve the name of the best embedder from the retriever node.

        Returns:
            Configuration of the best embedder.
        """
        best_retriever_artifact: EmbeddingArtifact = self._get_best_artifact(node_type=NodeType.embedding)  # type: ignore[assignment]
        return best_retriever_artifact.config

    def get_best_train_scores(self) -> NDArray[np.float64] | None:
        """Retrieve the train scores from the best scorer node.

        Returns:
            Train scores as a numpy array.
        """
        best_scorer_artifact: ScorerArtifact = self._get_best_artifact(node_type=NodeType.scoring)  # type: ignore[assignment]
        return best_scorer_artifact.train_scores

    def get_best_validation_scores(self) -> NDArray[np.float64] | None:
        """Retrieve the validation scores from the best scorer node.

        Returns:
            Validation scores as a numpy array.
        """
        best_scorer_artifact: ScorerArtifact = self._get_best_artifact(node_type=NodeType.scoring)  # type: ignore[assignment]
        return best_scorer_artifact.validation_scores

    def get_best_folded_scores(self) -> list[NDArray[np.float64]] | None:
        """Retrieve the validation scores from the best scorer node.

        Returns:
            Validation scores as a numpy array.
        """
        best_scorer_artifact: ScorerArtifact = self._get_best_artifact(node_type=NodeType.scoring)  # type: ignore[assignment]
        return best_scorer_artifact.folded_scores

    def get_best_test_scores(self) -> NDArray[np.float64] | None:
        """Retrieve the test scores from the best scorer node.

        Returns:
            Test scores as a numpy array.
        """
        best_scorer_artifact: ScorerArtifact = self._get_best_artifact(node_type=NodeType.scoring)  # type: ignore[assignment]
        return best_scorer_artifact.test_scores

    def dump_evaluation_results(self) -> dict[str, Any]:
        """Dump evaluation results for all nodes.

        Returns:
            Dictionary containing metrics and configurations for all nodes.
        """
        node_wise_metrics = {node_type: self._get_metrics_values(node_type) for node_type in NodeType}
        return {
            "pipeline_metrics": self.pipeline_metrics,
            "metrics": node_wise_metrics,
            "configs": self.trials.model_dump(),
        }

    def get_inference_nodes_config(self, asdict: bool = False) -> list[InferenceNodeConfig]:
        """Generate configuration for inference nodes based on the best trials.

        Args:
            asdict: Whether to return the configuration as dictionaries.

        Returns:
            List of configurations for inference nodes.
        """
        trial_ids = [self._get_best_trial_idx(node_type) for node_type in NodeType]
        res = []
        for idx, node_type in zip(trial_ids, NodeType, strict=True):
            if idx is None:
                continue
            trial = self.trials.get_trial(node_type, idx)
            item = {
                "node_type": node_type.value,
                "module_name": trial.module_name,
                "module_config": trial.module_params,
                "load_path": trial.module_dump_dir,
            }
            res.append(item if asdict else InferenceNodeConfig(**item))  # type: ignore[arg-type]
        return res  # type: ignore[return-value]

    def _get_best_module(self, node_type: str) -> "BaseModule | None":
        """Retrieve the best module for a specific node type.

        Args:
            node_type: Type of the node.

        Returns:
            The best module, or None if no best trial exists.
        """
        idx = self._get_best_trial_idx(node_type)
        if idx is not None:
            return self.modules.get(node_type)[idx]
        return None

    def get_best_modules(self) -> dict[NodeType, "BaseModule"]:
        """Retrieve the best modules for all node types.

        Returns:
            Dictionary of the best modules for each node type.
        """
        res = {nt: self._get_best_module(nt) for nt in NodeType}
        return {nt: m for nt, m in res.items() if m is not None}
