"""Custom ExhaustiveCategoricalSampler for exhaustive categorical search space exploration.

Extends Optuna's BruteForceSampler to:
1. Enforce categorical-only parameters (reject Float/Int distributions)
2. Raise SearchSpaceExhausted instead of silently duplicating combinations
3. Track total_combinations for operator-side excess Pod prevention
"""

import itertools
from functools import reduce
from operator import mul
from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence

from optuna.distributions import BaseDistribution, CategoricalDistribution
from optuna.samplers import BruteForceSampler
from optuna.trial import FrozenTrial, TrialState

if TYPE_CHECKING:
    from optuna.study import Study


class SearchSpaceExhausted(Exception):  # noqa: N818
    """Raised when all categorical combinations have been exhausted."""

    pass


class ExhaustiveCategoricalSampler(BruteForceSampler):
    """Sampler that exhaustively explores all categorical parameter combinations exactly once.

    Unlike BruteForceSampler which falls back to random sampling when combinations
    are exhausted, this sampler raises SearchSpaceExhausted to prevent wasted trials.

    Uses sample_relative() with itertools.product to bypass _TreeNode, avoiding
    param_name mismatch crash in distributed gRPC storage (protobuf map ordering).

    Search space is auto-discovered from the first completed trial's distributions.
    First-batch trials (before any trial completes) use sample_independent() with
    simple random selection.

    Args:
        seed: Random seed for sampling order. Not recommended for distributed settings.
    """

    def __init__(self, seed: Optional[int] = None) -> None:
        super().__init__(seed=seed)

    def infer_relative_search_space(
        self,
        study: "Study",
        trial: FrozenTrial,
    ) -> Dict[str, BaseDistribution]:
        """Auto-discover search space from completed trials' distributions.

        Returns empty dict when no completed trials exist (first-batch),
        causing all params to route to sample_independent() for random fallback.
        """
        trials = study._storage.get_all_trials(
            study._study_id,
            deepcopy=False,
            states=(TrialState.COMPLETE, TrialState.PRUNED),
        )
        if not trials:
            return {}
        return dict(sorted(trials[0].distributions.items()))

    def sample_relative(
        self,
        study: "Study",
        trial: FrozenTrial,
        search_space: Dict[str, BaseDistribution],
    ) -> Dict[str, Any]:
        """Sample by enumerating all combinations and picking an unused one.

        Bypasses _TreeNode entirely to avoid dict ordering crash.
        Uses fixed sorted key ordering and set-based deduplication.
        """
        if not search_space:
            return {}

        param_names = sorted(search_space.keys())

        trials = study._storage.get_all_trials(
            study._study_id,
            deepcopy=False,
            states=(
                TrialState.COMPLETE,
                TrialState.PRUNED,
                TrialState.RUNNING,
                TrialState.FAIL,
            ),
        )
        used = set()
        for t in trials:
            if t.number == trial.number:
                continue
            combo = tuple(t.params.get(name) for name in param_names)
            if None not in combo:
                used.add(combo)

        all_values = [list(search_space[name].choices) for name in param_names]
        all_combos = list(itertools.product(*all_values))
        unused = [c for c in all_combos if c not in used]

        if not unused:
            raise SearchSpaceExhausted("All categorical combinations have been exhausted")

        idx = self._rng.rng.randint(len(unused))
        return dict(zip(param_names, unused[idx]))

    def sample_independent(
        self,
        study: "Study",
        trial: FrozenTrial,
        name: str,
        param_distribution: BaseDistribution,
    ) -> Any:
        """First-batch fallback: random selection for trials before any completion.

        Only reached when infer_relative_search_space() returns {} (no completed trials).
        Also serves as a guard rejecting non-CategoricalDistribution.
        """
        if not isinstance(param_distribution, CategoricalDistribution):
            raise ValueError(
                f"ExhaustiveCategoricalSampler only supports CategoricalDistribution, "
                f"got {type(param_distribution).__name__} for parameter '{name}'"
            )
        idx = self._rng.rng.randint(len(param_distribution.choices))
        return param_distribution.choices[idx]

    def after_trial(
        self,
        study: "Study",
        trial: FrozenTrial,
        state: TrialState,
        values: Optional[Sequence[float]],
    ) -> None:
        """Set _total_combinations and stop study when all combinations exhausted.

        Does NOT call super().after_trial() to avoid _TreeNode crash.
        Uses set-based unique combination counting instead.
        """
        if not trial.distributions:
            return

        param_names = sorted(trial.distributions.keys())

        tc = reduce(
            mul,
            (
                len(d.choices)
                for d in trial.distributions.values()
                if isinstance(d, CategoricalDistribution)
            ),
            1,
        )
        study.set_user_attr("_total_combinations", str(tc))

        # FAIL trials don't contribute to exhaustion count; skip storage query
        if state not in (TrialState.COMPLETE, TrialState.PRUNED):
            return

        # Optuna calls after_trial BEFORE set_trial_state_values, so the current
        # trial is still RUNNING in storage. Query past completed trials and manually
        # include the current trial's combo using the state parameter.
        trials = study._storage.get_all_trials(
            study._study_id,
            deepcopy=False,
            states=(TrialState.COMPLETE, TrialState.PRUNED),
        )
        completed = set()
        for t in trials:
            combo = tuple(t.params.get(name) for name in param_names)
            if None not in combo:
                completed.add(combo)

        # Add current trial (still RUNNING in storage, but completing)
        current_combo = tuple(trial.params.get(name) for name in param_names)
        if None not in current_combo:
            completed.add(current_combo)

        if len(completed) >= tc:
            study.stop()
