from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Protocol

import numpy as np

from zen_fronts.policies import DefaultParentPolicy, ParentPolicy
from zen_fronts.selection import RefreshResult


class ParentSelector(Protocol):
    """Strategy object for `ZenFronts.choose_parent` semantics."""

    def choose_parent(
        self,
        loser_id: int,
        *,
        refresh_result: RefreshResult,
        rng: np.random.Generator,
        **kwargs: Any,
    ) -> int:
        raise NotImplementedError


@dataclass(frozen=True)
class DefaultParentSelector:
    """Default selector for `ZenFronts.choose_parent`.

    Behavior
    --------
    If `prefer_quality=True` and `refresh_result.quality_score` is available,
    chooses the parent as the winner with maximum `quality_score` (tie-break by id).

    Otherwise delegates to the provided :class:`~zen_fronts.policies.ParentPolicy`.
    """

    policy: ParentPolicy = DefaultParentPolicy(stochastic=False)
    prefer_quality: bool = True

    def choose_parent(
        self,
        loser_id: int,
        *,
        refresh_result: RefreshResult,
        rng: np.random.Generator,
        **kwargs: Any,
    ) -> int:
        if bool(self.prefer_quality) and (refresh_result.quality_score is not None):
            winners = list(refresh_result.winners_ids)
            if not winners:
                winners = list(refresh_result.active_ids)
            if not winners:
                raise ValueError("Cannot choose a parent: no active points")

            idxs = np.array([refresh_result.index_of(pid) for pid in winners], dtype=np.int64)
            q = refresh_result.quality_score[idxs].astype(np.float64, copy=False)
            ids = np.asarray(winners, dtype=np.int64)

            # max quality, tie-break by smallest id
            order = np.lexsort((ids, -q))
            return int(ids[int(order[0])])

        # Fallback: delegate to policy (policy currently ignores kwargs).
        return int(self.policy.choose_parent(int(loser_id), refresh_result=refresh_result, rng=rng))
