"""IRT engine — 1PL/Rasch with Bayesian theta update (SLICE-04, REQ-NFR-IRT-01).

P_success(theta, b) = logistic(theta - b) = 1 / (1 + exp(-(theta - b))).
update_theta uses a Gaussian-approximation Bayesian update (Kalman-like):
the posterior precision is the prior precision plus the Fisher information
P*(1-P), and the posterior mean shifts toward the outcome by the Kalman gain.

Cold-start (R-IRT-01): theta=0, sigma_sq=1; until >=5 observations, scenario
selection falls back to difficulty-based matching (difficulty closest to
round(theta + logit(target_p))).
"""

from __future__ import annotations

import math

from server.scenarios.library import ScenarioLibrary
from server.scenarios.schema import Scenario

COLD_START_MIN_OBSERVATIONS = 5
DEFAULT_THETA = 0.0
DEFAULT_SIGMA_SQ = 1.0


def _logit(p: float) -> float:
    return math.log(p / (1.0 - p))


class IRTEngine:
    """1PL/Rasch IRT with Gaussian-approximation Bayesian theta updates."""

    @staticmethod
    def P_success(theta: float, b: float) -> float:
        exp_neg = math.exp(-(theta - b))
        return 1.0 / (1.0 + exp_neg)

    @staticmethod
    def update_theta(
        theta: float, sigma_sq: float, outcome: float, b: float
    ) -> tuple[float, float]:
        """Bayesian update of theta given a binary (0/1) outcome.

        Uses the standard 1PL Gaussian-approximation (Kalman-like) update:
            P = P_success(theta, b)
            new_precision = 1/sigma_sq + P*(1-P)
            new_sigma_sq = 1 / new_precision
            new_theta = theta + new_sigma_sq * (outcome - P)
        """
        p = IRTEngine.P_success(theta, b)
        prior_precision = 1.0 / sigma_sq
        info = p * (1.0 - p)
        new_precision = prior_precision + info
        new_sigma_sq = 1.0 / new_precision
        new_theta = theta + new_sigma_sq * (outcome - p)
        return new_theta, new_sigma_sq

    @staticmethod
    def select_scenario(
        theta: float,
        library: ScenarioLibrary,
        path: str,
        target_p: float = 0.7,
        observations: int = 0,
    ) -> Scenario | None:
        """Select the next scenario for a learner.

        If observations < COLD_START_MIN_OBSERVATIONS (R-IRT-01), fall back to
        difficulty-based selection: pick the scenario whose `difficulty` is
        closest to round(theta + logit(target_p)).

        Otherwise delegate to library.select_for_theta (IRT-aware selection
        targeting ~target_p).
        """
        if observations < COLD_START_MIN_OBSERVATIONS:
            entries = library.list_by_path(path)
            if not entries:
                return None
            target_difficulty = round(theta + _logit(target_p))
            target_difficulty = max(1, min(5, target_difficulty))
            best_entry = None
            best_dist = math.inf
            for e in entries:
                dist = abs(e.difficulty - target_difficulty)
                if dist < best_dist:
                    best_dist = dist
                    best_entry = e
            if best_entry is None:
                return None
            return library.get(best_entry.id)
        return library.select_for_theta(theta, path, target_p=target_p)


__all__ = [
    "IRTEngine",
    "COLD_START_MIN_OBSERVATIONS",
    "DEFAULT_THETA",
    "DEFAULT_SIGMA_SQ",
]