"""High-level facade: ZenFronts.

This module ties together:

* :class:`zen_fronts.crits.CriteriaStore` (async objective storage)
* :class:`zen_fronts.selection.SelectionEngine` (Monte-Carlo winner/loser selection)
* :class:`zen_fronts.runtime.ParentSelector` (parent selection semantics)
* :class:`zen_fronts.runtime.Performant` (new individual creation semantics)

The facade intentionally keeps parameter sampling/mutation injectable so that
`ParamSpace` can evolve independently.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Mapping

import numpy as np

from zen_fronts.crits import CriteriaStore, CriteriaStoreConfig
from zen_fronts.policies import DefaultParentPolicy, ParentPolicy
from zen_fronts.runtime import DefaultParentSelector, DefaultPerformant, MutatorFn, ParentSelector, Performant
from zen_fronts.selection import RefreshResult, SelectionConfig, SelectionEngine
from zen_fronts.selection.schema import SCHEMA_NAME as SELECTION_SCHEMA_NAME
from zen_fronts.selection.schema import SCHEMA_VERSION as SELECTION_SCHEMA_VERSION


SamplerFn = Callable[[np.random.Generator], dict[str, Any]]


@dataclass(frozen=True)
class ZenFrontsConfig:
    seed: int = 0


class ZenFronts:
    """A minimal but usable orchestrator for async multi-objective PBT.

    Key invariants
    --------------
    * `CriteriaStore` is the source of truth for objective statistics.
    * `SelectionEngine.refresh()` produces a `RefreshResult` and the facade
      persists per-individual selection summaries for downstream consumption.
    * Parent selection and creation of new individuals are delegated to small
      strategy objects (`ParentSelector` and `Performant`).
    """

    def __init__(
        self,
        *,
        crits: Mapping[str, str],
        selection: SelectionConfig | None = None,
        # Backwards-compatible: still allow a policy; used by default selector.
        parent_policy: ParentPolicy | None = None,
        parent_selector: ParentSelector | None = None,
        performant: Performant | None = None,
        crit_store_config: CriteriaStoreConfig | None = None,
        sampler: SamplerFn | None = None,
        mutator: MutatorFn | None = None,
        config: ZenFrontsConfig | None = None,
    ) -> None:
        self._cfg = config or ZenFrontsConfig()
        self._rng = np.random.default_rng(int(self._cfg.seed))

        self._crits = CriteriaStore(crits, config=crit_store_config)
        self._selector = SelectionEngine(self._crits, config=selection)

        # Strategy objects
        if parent_selector is not None:
            self._parent_selector = parent_selector
        else:
            self._parent_selector = DefaultParentSelector(
                policy=parent_policy or DefaultParentPolicy(stochastic=False)
            )

        self._performant = performant or DefaultPerformant()

        self._sampler = sampler
        self._mutator = mutator

        self._next_id: int = 0
        self._params: dict[int, dict[str, Any]] = {}

        # Per-individual selection summaries persisted after refresh().
        self._selection_epoch: int = 0
        self._selection_stats: dict[int, dict[str, Any]] = {}

    # ------------------------------ points API ------------------------------

    @property
    def points(self) -> dict[int, dict[str, Any]]:
        """Mapping of active point_id -> params."""

        return self._params

    def add_point(self, params: Mapping[str, Any]) -> int:
        pid = self._next_id
        self._next_id += 1

        self._params[pid] = dict(params)
        self._crits.register_point(pid, active=True)
        return pid

    def add_random_point(self, n: int = 1) -> list[int]:
        if self._sampler is None:
            raise RuntimeError("ZenFronts.add_random_point requires `sampler=`")
        ids: list[int] = []
        for _ in range(int(n)):
            ids.append(self.add_point(self._sampler(self._rng)))
        return ids

    def delete_point(self, point_id: int) -> None:
        pid = int(point_id)
        if pid in self._params:
            del self._params[pid]
        # Tombstone for selection: point stops being active, but objective stats remain.
        self._crits.set_active(pid, False)

    def params(self, point_id: int) -> dict[str, Any]:
        pid = int(point_id)
        if pid not in self._params:
            raise KeyError(pid)
        return dict(self._params[pid])

    def info(self, point_id: int) -> dict[str, Any]:
        pid = int(point_id)
        st = {
            spec.name: (
                None
                if self._crits.state(pid, spec.name) is None
                else {
                    "mu": float(self._crits.state(pid, spec.name).mu),
                    "trend": float(self._crits.state(pid, spec.name).trend),
                    "sigma2": float(self._crits.state(pid, spec.name).sigma2),
                    "t_last": float(self._crits.state(pid, spec.name).t_last),
                    "ready": bool(self._crits.state(pid, spec.name).ready),
                    "source": str(self._crits.state(pid, spec.name).source),
                }
            )
            for spec in self._crits.crit_specs
        }
        return {
            "point_id": pid,
            "active": pid in self._params,
            "params": self._params.get(pid),
            "crits": st,
            # New: persisted selection summaries after the latest refresh() that included this point.
            "selection": self._selection_stats.get(pid),
        }

    # ------------------------------ criteria updates ------------------------------

    def update_crits(
        self,
        point_id: int,
        raw_or_stats: Mapping[str, float | Mapping[str, float]],
        *,
        t: float | None = None,
    ) -> None:
        self._crits.update_crits(point_id, raw_or_stats, t=t)

    # ------------------------------ selection cycle ------------------------------

    def _persist_refresh_stats(self, rr: RefreshResult) -> None:
        self._selection_epoch += 1
        epoch = self._selection_epoch

        n_active = len(rr.active_ids)
        # k = ceil(percentile * N), clamped to [1..N] when N>0.
        if n_active > 0:
            k = int(np.ceil(float(rr.percentile) * float(n_active)))
            k = max(1, min(k, n_active))
        else:
            k = 0

        winners = set(rr.winners_ids)
        losers = set(rr.losers_ids)

        # If detailed stats are unavailable, still persist aggregate counts.
        for pid in rr.active_ids:
            i = rr.index_of(pid)
            point_stats: dict[str, Any] = {
                "schema_name": SELECTION_SCHEMA_NAME,
                "schema_version": SELECTION_SCHEMA_VERSION,
                "point_id": int(pid),
                "epoch": int(epoch),
                "n_samples": int(rr.n_samples),
                "percentile": float(rr.percentile),
                "seed_used": int(rr.seed_used),
                "n_active": int(n_active),
                "k": int(k),
                "is_winner": pid in winners,
                "is_loser": pid in losers,
                "win_count": int(rr.win_count[i]),
                "lose_count": int(rr.lose_count[i]),
                "rank_sum": int(rr.rank_sum[i]),
                "quality_score": (None if rr.quality_score is None else float(rr.quality_score[i])),
                "quantiles_mode_used": (None if rr.quantiles_mode_used is None else int(rr.quantiles_mode_used)),
            }

            def _metric_block(prefix: str) -> dict[str, float] | None:
                mean = getattr(rr, f"{prefix}_mean")
                if mean is None:
                    return None
                med = getattr(rr, f"{prefix}_median")
                std = getattr(rr, f"{prefix}_std")
                q25 = getattr(rr, f"{prefix}_q25")
                q75 = getattr(rr, f"{prefix}_q75")
                return {
                    "mean": float(mean[i]),
                    "median": float(med[i]) if med is not None else float(mean[i]),
                    "std": float(std[i]) if std is not None else 0.0,
                    "q25": float(q25[i]) if q25 is not None else float(mean[i]),
                    "q75": float(q75[i]) if q75 is not None else float(mean[i]),
                }

            point_stats["place"] = _metric_block("place")
            point_stats["front"] = _metric_block("front")
            point_stats["within"] = _metric_block("within")

            self._selection_stats[int(pid)] = point_stats

    def refresh(self, *, now: float | None = None, seed: int | None = None) -> tuple[int, ...]:
        losers = self._selector.refresh(now=now, seed=seed)
        rr = self._selector.last_result
        if rr is not None:
            self._persist_refresh_stats(rr)
        return losers

    def choose_parent(self, loser_id: int, *, seed: int | None = None, **kwargs: Any) -> int:
        rr = self._selector.last_result
        if rr is None:
            raise RuntimeError("choose_parent() requires refresh() first")

        rng = np.random.default_rng(int(seed)) if seed is not None else self._rng
        ctx = {
            "zf_selection_epoch": int(self._selection_epoch),
            "zf_loser_id": int(loser_id),
            "zf_loser_selection": self._selection_stats.get(int(loser_id)),
        }
        ctx.update(kwargs)

        return int(
            self._parent_selector.choose_parent(
                int(loser_id),
                refresh_result=rr,
                rng=rng,
                **ctx,
            )
        )

    def perform_new(
        self,
        parent_id: int,
        *,
        looser: int | None = None,
        remove_looser: bool = True,
        seed: int | None = None,
        **kwargs: Any,
    ) -> tuple[int, dict[str, Any]]:
        if self._mutator is None:
            raise RuntimeError("ZenFronts.perform_new requires `mutator=`")

        pid = int(parent_id)
        rng = np.random.default_rng(int(seed)) if seed is not None else self._rng

        rr = self._selector.last_result

        ctx = {
            "zf_parent_id": pid,
            "zf_loser_id": int(looser) if looser is not None else None,
            "zf_remove_looser": bool(remove_looser),
            "zf_selection_epoch": int(self._selection_epoch),
            "zf_refresh_result": rr,
            "zf_parent_selection": self._selection_stats.get(pid),
            "zf_loser_selection": self._selection_stats.get(int(looser)) if looser is not None else None,
        }
        # User kwargs may include step counters, external metadata, etc.
        ctx.update(kwargs)

        return self._performant.perform_new(
            pid,
            get_params=self.params,
            add_point=self.add_point,
            delete_point=self.delete_point,
            rng=rng,
            mutator=self._mutator,
            looser_id=int(looser) if looser is not None else None,
            remove_looser=bool(remove_looser),
            **ctx,
        )
