"""
IRT 2PL (Two-Parameter Logistic) engine for Computerized Adaptive Testing (CAT).

Provides:
  - IRT2PLEngine.estimate_theta_eap(responses)  -> (theta, sem)
  - IRT2PLEngine.select_next_item(theta, candidate_items, administered_ids) -> item | None

The EAP (Expected A Posteriori) estimator uses a standard-normal prior and
41-point Gauss-Legendre quadrature over [-4, +4].  Item selection uses
Maximum Fisher Information, which for the 2PL model is maximised when the
item difficulty b is close to the current ability estimate theta.
"""
from functools import lru_cache
import math


class IRT2PLEngine:
    """Core IRT 2PL engine: ability estimation + adaptive item selection."""

    # 41-point quadrature nodes over [-4, +4]
    QUAD_NODES = [-4.0 + i * (8.0 / 40.0) for i in range(41)]

    # ------------------------------------------------------------------ #
    #  Probability helpers
    # ------------------------------------------------------------------ #
    @staticmethod
    @lru_cache(maxsize=4096)
    def _probability(theta: float, a: float, b: float) -> float:
        """P(correct | theta, a, b) under the 2PL model."""
        return 1.0 / (1.0 + math.exp(-a * (theta - b)))

    @staticmethod
    def probability(theta: float, a: float, b: float) -> float:
        """Public wrapper around the cached probability."""
        return IRT2PLEngine._probability(round(theta, 4), round(a, 4), round(b, 4))

    # ------------------------------------------------------------------ #
    #  Ability estimation  (EAP with standard-normal prior)
    # ------------------------------------------------------------------ #
    @staticmethod
    def estimate_theta_eap(responses) -> tuple:
        """
        Estimate ability (theta) and standard error of measurement (SEM)
        via Expected A Posteriori estimation.

        Parameters
        ----------
        responses : list[tuple]
            Each tuple is (a, b, is_correct) where
              a = discrimination parameter
              b = difficulty parameter
              is_correct = bool whether the response was correct

        Returns
        -------
        (theta, sem) : tuple[float, float]
        """
        if not responses:
            return 0.0, 1.0

        nodes = IRT2PLEngine.QUAD_NODES
        weights = []
        for x in nodes:
            prior = (1.0 / math.sqrt(2 * math.pi)) * math.exp(-0.5 * (x ** 2))
            likelihood = 1.0
            for a, b, is_correct in responses:
                p = IRT2PLEngine._probability(round(x, 4), round(a, 4), round(b, 4))
                likelihood *= p if is_correct else (1.0 - p)
            weights.append(prior * likelihood)

        sum_w = sum(weights)
        if sum_w == 0:
            return 0.0, 1.0

        norm_weights = [w / sum_w for w in weights]
        eap_theta = sum(w * x for w, x in zip(norm_weights, nodes))
        variance = sum(
            w * ((x - eap_theta) ** 2) for w, x in zip(norm_weights, nodes)
        )

        return round(eap_theta, 4), round(math.sqrt(variance), 4)

    # ------------------------------------------------------------------ #
    #  Adaptive item selection  (Maximum Fisher Information)
    # ------------------------------------------------------------------ #
    @staticmethod
    def fisher_information(theta: float, a: float, b: float) -> float:
        """
        Fisher information for a single 2PL item at ability theta.

        I(theta) = a^2 * P * (1 - P)
        """
        p = IRT2PLEngine.probability(theta, a, b)
        return (a ** 2) * p * (1.0 - p)

    @staticmethod
    def select_next_item(
        theta: float = None,
        candidate_items=None,
        administered_ids=None,
        *,
        current_theta: float = None,
        administered_item_ids=None,
    ):
        """
        Select the next item that maximises Fisher Information at the
        current ability estimate, excluding already-administered items.

        Accepts both ``theta``/``administered_ids`` and the aliases
        ``current_theta``/``administered_item_ids`` for backward compatibility.

        Parameters
        ----------
        theta / current_theta : float
            Current ability estimate for the dimension.
        candidate_items : list
            ORM AssessmentItem objects with .id, .a_parameter, .b_parameter.
        administered_ids / administered_item_ids : set
            IDs of items already administered in this session/dimension.

        Returns
        -------
        The AssessmentItem with the highest Fisher Information, or None if
        no candidates remain.
        """
        # Resolve aliases
        if theta is None:
            theta = current_theta
        if administered_ids is None:
            administered_ids = administered_item_ids
        if theta is None:
            theta = 0.0
        if administered_ids is None:
            administered_ids = set()
        if candidate_items is None:
            candidate_items = []

        best_item = None
        best_info = -1.0

        for item in candidate_items:
            if item.id in administered_ids:
                continue
            a = getattr(item, "a_parameter", None)
            b = getattr(item, "b_parameter", None)
            if a is None or b is None:
                continue
            info = IRT2PLEngine.fisher_information(theta, a, b)
            if info > best_info:
                best_info = info
                best_item = item

        return best_item


# Backward-compatible alias (some modules import OptimizedIRT2PLEngine)
OptimizedIRT2PLEngine = IRT2PLEngine