"""Selection engine on top of Monte-Carlo Pareto ranking.

Responsibilities
----------------
* Obtain active point ids from :class:`~zen_fronts.crits.CriteriaStore`.
* Build (mu, sigma2) matrices for those points.
* Run Monte-Carlo simulations to estimate winner/loser likelihood.
* Map winners/losers back to point ids and store a :class:`RefreshResult`.

The engine itself is intentionally policy-free: it does not decide how to pick
parents or how to mutate parameters. That is handled by parent selectors and
performants at the `ZenFronts` facade level.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from zen_fronts.crits.store import CriteriaStore
from zen_fronts.selection_core import mc_rank_select, mc_rank_select_detailed


@dataclass(frozen=True)
class SelectionConfig:
    n_samples: int = 256
    percentile: float = 0.2
    seed: int | None = 0
    mode_i: int = 1
    within_mode_i: int = 0

    # Quantiles strategy for detailed stats:
    #   0 = exact (store all samples; heavier memory)
    #   1 = streaming P² (constant memory; approximate quantiles)
    #   2 = auto (default; exact for small N*n_samples, else streaming)
    quantiles_mode_i: int = 2

    # Auto-quantiles budget for quantiles_mode_i=2.
    # Interpreted as a threshold on (N * n_samples).
    quantiles_budget: int = 2_500_000

    # Composite per-point quality score: quality = quality_sign * (
    #   w_place*place_mean + w_front*front_mean + w_within*within_mean
    # )
    # By default quality_sign=-1, so larger is better (best individual tends to 0).
    quality_w_place: float = 1.0
    quality_w_front: float = 1.0
    quality_w_within: float = 1.0
    quality_sign: float = -1.0

    # New: store per-point MC summaries (place/front/within).
    collect_stats: bool = True


@dataclass(frozen=True)
class RefreshResult:
    """Immutable snapshot of one `refresh()` call."""

    active_ids: tuple[int, ...]
    winners_ids: tuple[int, ...]
    losers_ids: tuple[int, ...]

    win_count: np.ndarray  # (N,) int32 aligned with active_ids
    lose_count: np.ndarray  # (N,) int32 aligned with active_ids
    rank_sum: np.ndarray  # (N,) int64 aligned with active_ids

    # Per-point summaries across MC samples (aligned with active_ids).
    place_mean: np.ndarray | None = None
    place_median: np.ndarray | None = None
    place_std: np.ndarray | None = None
    place_q25: np.ndarray | None = None
    place_q75: np.ndarray | None = None

    front_mean: np.ndarray | None = None
    front_median: np.ndarray | None = None
    front_std: np.ndarray | None = None
    front_q25: np.ndarray | None = None
    front_q75: np.ndarray | None = None

    within_mean: np.ndarray | None = None
    within_median: np.ndarray | None = None
    within_std: np.ndarray | None = None
    within_q25: np.ndarray | None = None
    within_q75: np.ndarray | None = None

    # Composite quality score aligned with active_ids (larger is better by default).
    quality_score: np.ndarray | None = None

    # For auto-quantiles: indicates which quantiles mode was actually used.
    # 0 = exact, 1 = streaming (P²), None = unknown/unset.
    quantiles_mode_used: int | None = None

    n_samples: int = 0
    percentile: float = 0.2
    seed_used: int = 0

    def index_of(self, point_id: int) -> int:
        return self.active_ids.index(int(point_id))

    def stats_of(self, point_id: int) -> dict[str, float | int | None]:
        i = self.index_of(point_id)
        out: dict[str, float | int | None] = {
            "win_count": int(self.win_count[i]),
            "lose_count": int(self.lose_count[i]),
            "rank_sum": int(self.rank_sum[i]),
        }

        if self.place_mean is not None:
            out.update(
                {
                    "quality_score": float(self.quality_score[i]) if self.quality_score is not None else None,
                    "place_mean": float(self.place_mean[i]),
                    "place_median": float(self.place_median[i]) if self.place_median is not None else None,
                    "place_std": float(self.place_std[i]) if self.place_std is not None else None,
                    "front_mean": float(self.front_mean[i]) if self.front_mean is not None else None,
                    "front_median": float(self.front_median[i]) if self.front_median is not None else None,
                    "front_std": float(self.front_std[i]) if self.front_std is not None else None,
                    "within_mean": float(self.within_mean[i]) if self.within_mean is not None else None,
                    "within_median": float(self.within_median[i]) if self.within_median is not None else None,
                    "within_std": float(self.within_std[i]) if self.within_std is not None else None,
                }
            )
        return out


class SelectionEngine:
    def __init__(self, store: CriteriaStore, *, config: SelectionConfig | None = None):
        self._store = store
        self._config = config or SelectionConfig()
        self._last: RefreshResult | None = None

    @property
    def last_result(self) -> RefreshResult | None:
        return self._last

    def refresh(self, *, now: float | None = None, seed: int | None = None) -> tuple[int, ...]:
        """Run one selection step.

        Returns a tuple of loser point ids (may be empty).
        """

        active_ids = tuple(self._store.active_point_ids())
        seed_used = (
            int(seed)
            if seed is not None
            else int(self._config.seed)
            if self._config.seed is not None
            else 0
        )

        if len(active_ids) == 0:
            self._last = RefreshResult(
                active_ids=(),
                winners_ids=(),
                losers_ids=(),
                win_count=np.zeros(0, dtype=np.int32),
                lose_count=np.zeros(0, dtype=np.int32),
                rank_sum=np.zeros(0, dtype=np.int64),
                n_samples=int(self._config.n_samples),
                percentile=float(self._config.percentile),
                seed_used=seed_used,
            )
            return ()

        # With a single active point, selection is undefined; treat as no-op.
        if len(active_ids) == 1:
            self._last = RefreshResult(
                active_ids=active_ids,
                winners_ids=active_ids,
                losers_ids=(),
                win_count=np.array([self._config.n_samples], dtype=np.int32),
                lose_count=np.array([0], dtype=np.int32),
                rank_sum=np.array([0], dtype=np.int64),
                place_mean=np.array([0.0], dtype=np.float64),
                place_median=np.array([0.0], dtype=np.float64),
                place_std=np.array([0.0], dtype=np.float64),
                front_mean=np.array([0.0], dtype=np.float64),
                front_median=np.array([0.0], dtype=np.float64),
                front_std=np.array([0.0], dtype=np.float64),
                within_mean=np.array([0.0], dtype=np.float64),
                within_median=np.array([0.0], dtype=np.float64),
                within_std=np.array([0.0], dtype=np.float64),
                quality_score=np.array([0.0], dtype=np.float64),
                quantiles_mode_used=0,
                n_samples=int(self._config.n_samples),
                percentile=float(self._config.percentile),
                seed_used=seed_used,
            )
            return ()

        mu, sigma2 = self._store.build_matrices(active_ids, now=now)

        if bool(self._config.collect_stats):
            (
                winners_i,
                losers_i,
                win_count,
                lose_count,
                rank_sum,
                place_stats,
                front_stats,
                within_stats,
                quality_score,
            ) = mc_rank_select_detailed(
                mu,
                sigma2,
                int(self._config.n_samples),
                percentile=float(self._config.percentile),
                seed=seed_used,
                mode_i=int(self._config.mode_i),
                within_mode_i=int(self._config.within_mode_i),
                quantiles_mode_i=int(self._config.quantiles_mode_i),
                quantiles_budget=int(self._config.quantiles_budget),
                quality_w_place=float(self._config.quality_w_place),
                quality_w_front=float(self._config.quality_w_front),
                quality_w_within=float(self._config.quality_w_within),
                quality_sign=float(self._config.quality_sign),
            )

            winners_ids = tuple(active_ids[int(i)] for i in np.asarray(winners_i, dtype=np.int64))
            losers_ids = tuple(active_ids[int(i)] for i in np.asarray(losers_i, dtype=np.int64))

            self._last = RefreshResult(
                active_ids=active_ids,
                winners_ids=winners_ids,
                losers_ids=losers_ids,
                win_count=np.asarray(win_count, dtype=np.int32),
                lose_count=np.asarray(lose_count, dtype=np.int32),
                rank_sum=np.asarray(rank_sum, dtype=np.int64),
                place_mean=np.asarray(place_stats["mean"], dtype=np.float64),
                place_median=np.asarray(place_stats["median"], dtype=np.float64),
                place_std=np.asarray(place_stats["std"], dtype=np.float64),
                place_q25=np.asarray(place_stats["q25"], dtype=np.float64),
                place_q75=np.asarray(place_stats["q75"], dtype=np.float64),
                front_mean=np.asarray(front_stats["mean"], dtype=np.float64),
                front_median=np.asarray(front_stats["median"], dtype=np.float64),
                front_std=np.asarray(front_stats["std"], dtype=np.float64),
                front_q25=np.asarray(front_stats["q25"], dtype=np.float64),
                front_q75=np.asarray(front_stats["q75"], dtype=np.float64),
                within_mean=np.asarray(within_stats["mean"], dtype=np.float64),
                within_median=np.asarray(within_stats["median"], dtype=np.float64),
                within_std=np.asarray(within_stats["std"], dtype=np.float64),
                within_q25=np.asarray(within_stats["q25"], dtype=np.float64),
                within_q75=np.asarray(within_stats["q75"], dtype=np.float64),
                quality_score=np.asarray(quality_score, dtype=np.float64),
                quantiles_mode_used=int(place_stats.get("quantiles_mode_used", 0)) if isinstance(place_stats, dict) else None,
                n_samples=int(self._config.n_samples),
                percentile=float(self._config.percentile),
                seed_used=seed_used,
            )

            return losers_ids

        # Legacy: aggregate-only selection.
        winners_i, losers_i, win_count, lose_count, rank_sum = mc_rank_select(
            mu,
            sigma2,
            int(self._config.n_samples),
            percentile=float(self._config.percentile),
            seed=seed_used,
            mode_i=int(self._config.mode_i),
            within_mode_i=int(self._config.within_mode_i),
        )

        winners_ids = tuple(active_ids[int(i)] for i in np.asarray(winners_i, dtype=np.int64))
        losers_ids = tuple(active_ids[int(i)] for i in np.asarray(losers_i, dtype=np.int64))

        self._last = RefreshResult(
            active_ids=active_ids,
            winners_ids=winners_ids,
            losers_ids=losers_ids,
            win_count=np.asarray(win_count, dtype=np.int32),
            lose_count=np.asarray(lose_count, dtype=np.int32),
            rank_sum=np.asarray(rank_sum, dtype=np.int64),
            n_samples=int(self._config.n_samples),
            percentile=float(self._config.percentile),
            seed_used=seed_used,
        )

        return losers_ids
