"""
Hyperparameter Search Space Definitions for Optuna.

Defines search distributions for each optimizable parameter using
configurable bounds from SearchSpaceOverrides (Pydantic V2).

Search ranges respect the type constraints in core/config/types.py
and default to domain-expert values for medical imaging, but can be
fully customized via YAML through OptunaConfig.search_space_overrides.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Callable, Dict, Optional

import optuna

if TYPE_CHECKING:  # pragma: no cover
    from orchard.core.config.optuna_config import SearchSpaceOverrides


def _default_overrides() -> SearchSpaceOverrides:
    """Lazy import to avoid circular dependency at module level."""
    from orchard.core.config.optuna_config import SearchSpaceOverrides

    return SearchSpaceOverrides()


# SEARCH SPACE DEFINITIONS
class SearchSpaceRegistry:
    """
    Centralized registry of hyperparameter search distributions.

    Reads bounds from a SearchSpaceOverrides instance, enabling full
    YAML customization of search ranges without code changes.

    Each method returns a dict of {param_name: suggest_function} where
    suggest_function takes a Trial object and returns a sampled value.

    Args:
        overrides: Configurable search range bounds. Uses defaults if None.
    """

    def __init__(self, overrides: Optional[SearchSpaceOverrides] = None):
        self.ov = overrides if overrides is not None else _default_overrides()

    def get_optimization_space(self) -> Dict[str, Callable]:
        """
        Core optimization hyperparameters (learning rate, weight decay, etc.).

        Returns:
            Dict mapping parameter names to sampling functions
        """
        ov = self.ov
        return {
            "learning_rate": lambda trial: trial.suggest_float(
                "learning_rate",
                ov.learning_rate.low,
                ov.learning_rate.high,
                log=ov.learning_rate.log,
            ),
            "weight_decay": lambda trial: trial.suggest_float(
                "weight_decay",
                ov.weight_decay.low,
                ov.weight_decay.high,
                log=ov.weight_decay.log,
            ),
            "momentum": lambda trial: trial.suggest_float(
                "momentum",
                ov.momentum.low,
                ov.momentum.high,
            ),
            "min_lr": lambda trial: trial.suggest_float(
                "min_lr",
                ov.min_lr.low,
                ov.min_lr.high,
                log=ov.min_lr.log,
            ),
        }

    def get_regularization_space(self) -> Dict[str, Callable]:
        """
        Regularization strategies (mixup, label smoothing, dropout).

        Returns:
            Dict of regularization parameter samplers
        """
        ov = self.ov
        return {
            "mixup_alpha": lambda trial: trial.suggest_float(
                "mixup_alpha",
                ov.mixup_alpha.low,
                ov.mixup_alpha.high,
            ),
            "label_smoothing": lambda trial: trial.suggest_float(
                "label_smoothing",
                ov.label_smoothing.low,
                ov.label_smoothing.high,
            ),
            "dropout": lambda trial: trial.suggest_float(
                "dropout",
                ov.dropout.low,
                ov.dropout.high,
            ),
        }

    def get_batch_size_space(self, resolution: int = 28) -> Dict[str, Callable]:
        """
        Batch size as categorical (resolution-aware).

        Args:
            resolution: Input image resolution (28 or 224)

        Returns:
            Dict with batch_size sampler
        """
        if resolution >= 224:
            batch_choices = list(self.ov.batch_size_high_res)
        else:
            batch_choices = list(self.ov.batch_size_low_res)

        return {
            "batch_size": lambda trial: trial.suggest_categorical("batch_size", batch_choices),
        }

    def get_scheduler_space(self) -> Dict[str, Callable]:
        """
        Learning rate scheduler parameters.

        Returns:
            Dict of scheduler-related samplers
        """
        ov = self.ov
        return {
            "cosine_fraction": lambda trial: trial.suggest_float(
                "cosine_fraction",
                ov.cosine_fraction.low,
                ov.cosine_fraction.high,
            ),
            "scheduler_patience": lambda trial: trial.suggest_int(
                "scheduler_patience",
                ov.scheduler_patience.low,
                ov.scheduler_patience.high,
            ),
        }

    def get_augmentation_space(self) -> Dict[str, Callable]:
        """
        Data augmentation intensity parameters.

        Returns:
            Dict of augmentation samplers
        """
        ov = self.ov
        return {
            "rotation_angle": lambda trial: trial.suggest_int(
                "rotation_angle",
                ov.rotation_angle.low,
                ov.rotation_angle.high,
            ),
            "jitter_val": lambda trial: trial.suggest_float(
                "jitter_val",
                ov.jitter_val.low,
                ov.jitter_val.high,
            ),
            "min_scale": lambda trial: trial.suggest_float(
                "min_scale",
                ov.min_scale.low,
                ov.min_scale.high,
            ),
        }

    def get_full_space(self, resolution: int = 28) -> Dict[str, Callable]:
        """
        Combined search space with all available parameters.

        Args:
            resolution: Input image resolution for batch size calculation

        Returns:
            Unified dict of all parameter samplers
        """
        full_space: Dict[str, Callable] = {}
        full_space.update(self.get_optimization_space())
        full_space.update(self.get_regularization_space())
        full_space.update(self.get_batch_size_space(resolution))
        full_space.update(self.get_scheduler_space())
        full_space.update(self.get_augmentation_space())
        return full_space

    def get_quick_space(self, resolution: int = 28) -> Dict[str, Callable]:
        """
        Reduced search space for fast exploration (most impactful params).

        Focuses on:
            - Learning rate (most critical)
            - Weight decay
            - Batch size (resolution-aware)
            - Dropout

        Args:
            resolution: Input image resolution for batch size calculation

        Returns:
            Dict of high-impact parameter samplers
        """
        space: Dict[str, Callable] = {}
        space.update(self.get_optimization_space())
        space.update(
            {
                "batch_size": self.get_batch_size_space(resolution)["batch_size"],
                "dropout": self.get_regularization_space()["dropout"],
            }
        )
        return space

    @staticmethod
    def get_model_space_224() -> Dict[str, Callable]:
        """Search space for 224x224 architectures with weight variants."""
        return {
            "model_name": lambda trial: trial.suggest_categorical(
                "model_name", ["resnet_18", "efficientnet_b0", "vit_tiny", "convnext_tiny"]
            ),
            "weight_variant": lambda trial: (
                trial.suggest_categorical(
                    "weight_variant",
                    [
                        None,  # Default variant
                        "vit_tiny_patch16_224.augreg_in21k_ft_in1k",
                        "vit_tiny_patch16_224.augreg_in21k",
                    ],
                )
                if trial.params.get("model_name") == "vit_tiny"
                else None
            ),
        }

    @staticmethod
    def get_model_space_28() -> Dict[str, Callable]:
        """Search space for 28x28 architectures."""
        return {
            "model_name": lambda trial: trial.suggest_categorical(
                "model_name", ["resnet_18", "mini_cnn"]
            ),
        }

    def get_full_space_with_models(self, resolution: int = 28) -> Dict[str, Callable]:
        """
        Combined search space including model architecture selection.

        Args:
            resolution: Input image resolution (determines model choices)

        Returns:
            Unified dict of all parameter samplers + model selection
        """
        full_space = self.get_full_space(resolution)

        if resolution >= 224:
            full_space.update(self.get_model_space_224())
        else:
            full_space.update(self.get_model_space_28())

        return full_space


# PRESET CONFIGURATIONS
def get_search_space(
    preset: str = "quick",
    resolution: int = 28,
    include_models: bool = False,
    overrides: Optional[SearchSpaceOverrides] = None,
):
    """
    Factory function to retrieve a search space preset.

    Args:
        preset: Name of the preset ("quick", "full", etc.)
        resolution: Input image resolution (affects batch_size choices)
        include_models: If True, includes model architecture selection
        overrides: Configurable search range bounds (uses defaults if None)

    Returns:
        Dict of parameter samplers

    Raises:
        ValueError: If preset name not recognized
    """
    registry = SearchSpaceRegistry(overrides)

    if preset == "quick":
        space = registry.get_quick_space(resolution)
    elif preset == "full":
        space = registry.get_full_space(resolution)
    else:
        raise ValueError(
            f"Unknown preset '{preset}'. Available: quick, full, "
            f"optimization_only, regularization_only"
        )

    if include_models:
        if resolution >= 224:
            space.update(registry.get_model_space_224())
        else:
            space.update(registry.get_model_space_28())

    return space


class FullSearchSpace:
    """
    Resolution-aware full search space with dynamic batch_size constraints.

    Prevents OOM errors by limiting batch sizes based on input resolution.
    Reads bounds from SearchSpaceOverrides for configurability.
    """

    def __init__(
        self,
        resolution: int = 28,
        overrides: Optional[SearchSpaceOverrides] = None,
    ):
        """
        Initialize search space with resolution context.

        Args:
            resolution: Input image resolution (28, 224, etc.)
            overrides: Configurable search range bounds (uses defaults if None)
        """
        self.resolution = resolution
        self.ov = overrides if overrides is not None else _default_overrides()

    def sample_params(self, trial: optuna.Trial) -> dict:
        """
        Sample hyperparameters with resolution-aware constraints.

        Args:
            trial: Optuna trial object

        Returns:
            Dict of sampled hyperparameters
        """
        ov = self.ov

        if self.resolution >= 224:
            batch_choices = list(ov.batch_size_high_res)
        else:
            batch_choices = list(ov.batch_size_low_res)

        return {
            "learning_rate": trial.suggest_float(
                "learning_rate",
                ov.learning_rate.low,
                ov.learning_rate.high,
                log=ov.learning_rate.log,
            ),
            "weight_decay": trial.suggest_float(
                "weight_decay",
                ov.weight_decay.low,
                ov.weight_decay.high,
                log=ov.weight_decay.log,
            ),
            "momentum": trial.suggest_float(
                "momentum",
                ov.momentum.low,
                ov.momentum.high,
            ),
            "min_lr": trial.suggest_float(
                "min_lr",
                ov.min_lr.low,
                ov.min_lr.high,
                log=ov.min_lr.log,
            ),
            "mixup_alpha": trial.suggest_float(
                "mixup_alpha",
                ov.mixup_alpha.low,
                ov.mixup_alpha.high,
            ),
            "label_smoothing": trial.suggest_float(
                "label_smoothing",
                ov.label_smoothing.low,
                ov.label_smoothing.high,
            ),
            "dropout": trial.suggest_float(
                "dropout",
                ov.dropout.low,
                ov.dropout.high,
            ),
            "batch_size": trial.suggest_categorical("batch_size", batch_choices),
            "cosine_fraction": trial.suggest_float(
                "cosine_fraction",
                ov.cosine_fraction.low,
                ov.cosine_fraction.high,
            ),
            "scheduler_patience": trial.suggest_int(
                "scheduler_patience",
                ov.scheduler_patience.low,
                ov.scheduler_patience.high,
            ),
            "rotation_angle": trial.suggest_int(
                "rotation_angle",
                ov.rotation_angle.low,
                ov.rotation_angle.high,
            ),
            "jitter_val": trial.suggest_float(
                "jitter_val",
                ov.jitter_val.low,
                ov.jitter_val.high,
            ),
            "min_scale": trial.suggest_float(
                "min_scale",
                ov.min_scale.low,
                ov.min_scale.high,
            ),
        }
