"""Parent selection policy.

Selection determines which points are winners/losers; the parent policy chooses
the parent for a given loser.

This module defines:

* An interface :class:`ParentPolicy`.
* A deterministic default policy (useful for reproducible tests).
* An optional stochastic mode (weighted sampling by win_count) with deterministic
  seeding.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from zen_fronts.selection.engine import RefreshResult


class ParentPolicy:
    def choose_parent(
        self,
        loser_id: int,
        *,
        refresh_result: RefreshResult,
        rng: np.random.Generator,
    ) -> int:  # pragma: no cover
        raise NotImplementedError


@dataclass(frozen=True)
class DefaultParentPolicy(ParentPolicy):
    """Default parent policy.

    Modes
    -----
    deterministic:
        Choose the winner with maximum `win_count`, tie-break by minimum
        `rank_sum`, then by point id.

    stochastic:
        Weighted sampling among winners proportional to `win_count`.
        Ties and determinism are resolved by the provided RNG.
    """

    stochastic: bool = False

    def choose_parent(
        self,
        loser_id: int,
        *,
        refresh_result: RefreshResult,
        rng: np.random.Generator,
    ) -> int:
        if len(refresh_result.active_ids) == 0:
            raise ValueError("Cannot choose a parent: no active points")

        winners = list(refresh_result.winners_ids)
        if not winners:
            # Fallback: choose the best by rank_sum among active.
            # Smaller rank_sum means more often earlier in order.
            active = list(refresh_result.active_ids)
            idxs = [refresh_result.index_of(pid) for pid in active]
            rank_sum = refresh_result.rank_sum[idxs]
            best_i = int(np.lexsort((np.asarray(active, dtype=np.int64), rank_sum))[0])
            return int(active[best_i])

        if not self.stochastic:
            # Deterministic: max win_count, then min rank_sum, then id.
            idxs = np.array([refresh_result.index_of(pid) for pid in winners], dtype=np.int64)
            win = refresh_result.win_count[idxs].astype(np.int64, copy=False)
            rank = refresh_result.rank_sum[idxs].astype(np.int64, copy=False)
            ids = np.asarray(winners, dtype=np.int64)
            order = np.lexsort((ids, rank, -win))
            return int(ids[int(order[0])])

        # Stochastic: sample proportional to win_count (or uniform if all zeros).
        idxs = np.array([refresh_result.index_of(pid) for pid in winners], dtype=np.int64)
        w = refresh_result.win_count[idxs].astype(np.float64)
        if not np.isfinite(w).all() or float(w.sum()) <= 0.0:
            probs = np.full(w.shape[0], 1.0 / w.shape[0], dtype=np.float64)
        else:
            probs = w / w.sum()

        choice = int(rng.choice(np.asarray(winners, dtype=np.int64), p=probs))
        return choice
