"""Exact enumeration of absence and blocked/disruption states.

The forecasting model's daily capacity distribution is a finite mixture over

- absence states (2^n), and
- blocked/disruption states (K),

giving K * 2^n daily regimes.

The Monte Carlo estimators in the main forecaster are unbiased and typically
sufficient, but exact enumeration can be useful for:

- diagnostics (which regimes dominate the left tail?),
- variance reduction (stratified sampling per regime), and
- very small teams (n <= 10) where 2^n is trivial.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterator, Sequence

import numpy as np


@dataclass(frozen=True)
class AbsenceState:
    """A particular subset of available developers on a given day."""

    mask: np.ndarray
    probability: float


@dataclass(frozen=True)
class DailyRegime:
    """A (absence state, blocked-state) regime for a given day."""

    absence_mask: np.ndarray
    blocked_state: int
    probability: float


def enumerate_absence_states(q: Sequence[float]) -> list[AbsenceState]:
    """Enumerate all 2^n absence states for independent Bernoulli availability.

    Parameters
    ----------
    q:
        Availability probabilities q_i for each developer.

    Returns
    -------
    A list of AbsenceState, each with a boolean mask of available developers and
    its exact probability.
    """
    q_arr = np.asarray(q, dtype=float).reshape(-1)
    if q_arr.size == 0:
        raise ValueError("q must be non-empty")
    if np.any((q_arr < 0) | (q_arr > 1)):
        raise ValueError("q must be in [0,1]")

    n = int(q_arr.size)
    states: list[AbsenceState] = []

    # Bitmask enumeration: 0..2^n-1.
    for bitmask in range(1 << n):
        mask = np.array([(bitmask >> i) & 1 for i in range(n)], dtype=bool)
        # Probability = prod_i q_i^{mask_i} (1-q_i)^{1-mask_i}
        prob = float(
            np.prod(np.where(mask, q_arr, 1.0 - q_arr))
        )
        states.append(AbsenceState(mask=mask, probability=prob))

    return states


def enumerate_daily_regimes(
    q: Sequence[float],
    blocked_probabilities: Sequence[float],
) -> list[DailyRegime]:
    """Enumerate all daily regimes (absence x blocked state).

    Parameters
    ----------
    q:
        Availability probabilities for the day.
    blocked_probabilities:
        Probabilities pi_k for blocked states k=0..K-1.

    Returns
    -------
    List of DailyRegime objects.
    """
    pi = np.asarray(blocked_probabilities, dtype=float).reshape(-1)
    if pi.size == 0:
        raise ValueError("blocked_probabilities must be non-empty")
    if np.any(pi < 0):
        raise ValueError("blocked_probabilities must be non-negative")
    total = float(np.sum(pi))
    if not np.isclose(total, 1.0):
        raise ValueError("blocked_probabilities must sum to 1")

    absence_states = enumerate_absence_states(q)
    regimes: list[DailyRegime] = []
    for k, p_k in enumerate(pi):
        for state in absence_states:
            regimes.append(
                DailyRegime(
                    absence_mask=state.mask,
                    blocked_state=int(k),
                    probability=float(state.probability * p_k),
                )
            )
    return regimes


def iter_daily_regimes(
    q: Sequence[float],
    blocked_probabilities: Sequence[float],
) -> Iterator[DailyRegime]:
    """Generator form of :func:`enumerate_daily_regimes` to avoid large lists."""

    pi = np.asarray(blocked_probabilities, dtype=float).reshape(-1)
    total = float(np.sum(pi))
    if not np.isclose(total, 1.0):
        raise ValueError("blocked_probabilities must sum to 1")

    absence_states = enumerate_absence_states(q)
    for k, p_k in enumerate(pi):
        for state in absence_states:
            yield DailyRegime(
                absence_mask=state.mask,
                blocked_state=int(k),
                probability=float(state.probability * p_k),
            )
