"""Criteria store.

This is the core stateful component that makes `zen_fronts` practical in
asynchronous settings:

* Points can report criteria at different times.
* Some criteria can be delivered as *raw* observations (floats) that need an
  online filter to produce (mu, trend, sigma2).
* Other criteria can be delivered as *external stats* (mu/trend/sigma2), e.g.
  computed by the trainer.

The store keeps all states, and can build dense matrices (mu, sigma2) for an
arbitrary point subset, suitable for passing into the Monte-Carlo ranking core.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Any, Iterable, Mapping

import time

import numpy as np

from .owrf_adapter import make_owrf_model, project_mu


class CritDirection(str, Enum):
    """Objective direction."""

    MAX = "max"
    MIN = "min"


class CritSource(str, Enum):
    RAW_OWRF = "raw_owrf"
    EXTERNAL_STATS = "external_stats"


@dataclass(frozen=True)
class CritSpec:
    name: str
    direction: CritDirection = CritDirection.MAX

    @staticmethod
    def parse(name: str, direction: str | CritDirection) -> "CritSpec":
        if isinstance(direction, CritDirection):
            d = direction
        else:
            d = CritDirection(direction)
        return CritSpec(name=str(name), direction=d)


@dataclass
class CritState:
    """State for a specific (point_id, criterion)."""

    mu: float = 0.0
    trend: float = 0.0
    sigma2: float = float("inf")
    t_last: float = 0.0
    ready: bool = False
    source: CritSource = CritSource.EXTERNAL_STATS

    # Optional: only for RAW_OWRF source.
    owrf_model: Any | None = None


@dataclass(frozen=True)
class CriteriaStoreConfig:
    """Configuration for criteria matrix building."""

    project_trend: bool = True


class CriteriaStore:
    """Store of criterion states for multiple points."""

    def __init__(
        self,
        crits: Mapping[str, str | CritDirection] | Iterable[CritSpec],
        *,
        config: CriteriaStoreConfig | None = None,
    ) -> None:
        if isinstance(crits, Mapping):
            specs = [CritSpec.parse(k, v) for k, v in crits.items()]
        else:
            specs = list(crits)

        if not specs:
            raise ValueError("CriteriaStore requires at least one criterion")

        names = [s.name for s in specs]
        if len(names) != len(set(names)):
            raise ValueError(f"Duplicate criterion names: {names!r}")

        self._specs: tuple[CritSpec, ...] = tuple(specs)
        self._spec_by_name: dict[str, CritSpec] = {s.name: s for s in specs}
        self._config = config or CriteriaStoreConfig()

        # Active points participate in selection. Tombstoned points remain in
        # the store but do not appear in active lists.
        self._active_points: set[int] = set()

        # Nested states: point_id -> crit_name -> CritState
        self._states: dict[int, dict[str, CritState]] = {}

    # ------------------------------ point lifecycle ------------------------------

    def register_point(self, point_id: int, *, active: bool = True) -> None:
        point_id = int(point_id)
        self._states.setdefault(point_id, {})
        if active:
            self._active_points.add(point_id)
        else:
            self._active_points.discard(point_id)

    def set_active(self, point_id: int, active: bool) -> None:
        point_id = int(point_id)
        self._states.setdefault(point_id, {})
        if active:
            self._active_points.add(point_id)
        else:
            self._active_points.discard(point_id)

    # ------------------------------ state access ------------------------------

    @property
    def crit_specs(self) -> tuple[CritSpec, ...]:
        return self._specs

    def state(self, point_id: int, crit_name: str) -> CritState | None:
        d = self._states.get(int(point_id))
        if not d:
            return None
        return d.get(crit_name)

    def is_ready(self, point_id: int) -> bool:
        """True if the point has ready states for all criteria."""

        d = self._states.get(int(point_id))
        if not d:
            return False
        for spec in self._specs:
            st = d.get(spec.name)
            if st is None or not st.ready:
                return False
        return True

    def active_point_ids(self) -> list[int]:
        """Return active points that are ready on all criteria."""

        # Deterministic ordering is important for reproducible selection.
        ids = [pid for pid in self._active_points if self.is_ready(pid)]
        ids.sort()
        return ids

    # ------------------------------ update API ------------------------------

    def update_crits(
        self,
        point_id: int,
        raw_or_stats: Mapping[str, float | Mapping[str, float]],
        *,
        t: float | None = None,
    ) -> None:
        """Update one point's criteria.

        Parameters
        ----------
        point_id:
            Point identifier.
        raw_or_stats:
            Dict mapping criterion name -> either:
              * float (raw observation), or
              * mapping with keys: mu, trend, sigma2 (external stats).
        t:
            Timestamp for this update. If omitted, uses `time.time()`.
        """

        pid = int(point_id)
        self._states.setdefault(pid, {})

        if t is None:
            t = time.time()
        t = float(t)

        for name, payload in raw_or_stats.items():
            if name not in self._spec_by_name:
                raise KeyError(f"Unknown criterion: {name!r}")

            if isinstance(payload, Mapping):
                self._update_from_stats(pid, name, payload, t=t)
            else:
                self._update_from_raw(pid, name, float(payload), t=t)

    def _update_from_stats(
        self,
        point_id: int,
        name: str,
        stats: Mapping[str, float],
        *,
        t: float,
    ) -> None:
        mu = float(stats.get("mu"))
        trend = float(stats.get("trend", 0.0))
        sigma2 = float(stats.get("sigma2", 0.0))
        sigma2 = max(sigma2, 0.0)

        st = self._states[point_id].get(name)
        if st is None:
            st = CritState()
            self._states[point_id][name] = st

        st.mu = mu
        st.trend = trend
        st.sigma2 = sigma2
        st.t_last = float(t)
        st.ready = True
        st.source = CritSource.EXTERNAL_STATS
        st.owrf_model = None

    def _update_from_raw(self, point_id: int, name: str, y: float, *, t: float) -> None:
        st = self._states[point_id].get(name)
        if st is None:
            st = CritState()
            st.source = CritSource.RAW_OWRF
            st.owrf_model = make_owrf_model()
            self._states[point_id][name] = st

        if st.owrf_model is None:
            st.source = CritSource.RAW_OWRF
            st.owrf_model = make_owrf_model()

        st.owrf_model.update(float(y), float(t))
        mu, trend, sigma2, t_last = st.owrf_model.stats()
        st.mu = float(mu)
        st.trend = float(trend)
        st.sigma2 = max(float(sigma2), 0.0)
        st.t_last = float(t_last)
        st.ready = True
        st.source = CritSource.RAW_OWRF

    # ------------------------------ matrices for selection ------------------------------

    def build_matrices(
        self,
        point_ids: Iterable[int],
        *,
        now: float | None = None,
    ) -> tuple[np.ndarray, np.ndarray]:
        """Build dense mu/sigma2 matrices for selection.

        The resulting matrices have shape (N, M), where:

        * N == len(point_ids)
        * M == number of criteria

        Conventions:

        * Output is ALWAYS "maximize" (min objectives are multiplied by -1).
        * If configured and `now` is provided, ``mu`` is projected as
          ``mu + trend * (now - t_last)``.
        """

        ids = [int(x) for x in point_ids]
        specs = self._specs
        N = len(ids)
        M = len(specs)

        mu = np.empty((N, M), dtype=np.float64)
        sigma2 = np.empty((N, M), dtype=np.float64)

        if now is None:
            now = time.time()
        now = float(now)

        for i, pid in enumerate(ids):
            d = self._states.get(pid)
            if not d:
                raise KeyError(f"Unknown point_id: {pid}")

            for j, spec in enumerate(specs):
                st = d.get(spec.name)
                if st is None or not st.ready:
                    raise ValueError(
                        f"Point {pid} is not ready for criterion {spec.name!r}"
                    )

                mu_ij = float(st.mu)
                if self._config.project_trend and st.t_last and np.isfinite(st.t_last):
                    mu_ij = project_mu(mu_ij, float(st.trend), float(st.t_last), now)

                if spec.direction == CritDirection.MIN:
                    mu_ij = -mu_ij

                mu[i, j] = mu_ij
                sigma2[i, j] = max(float(st.sigma2), 0.0)

        return mu, sigma2
