"""Capacity (supply) simulation for the sprint forecasting model."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Tuple

import numpy as np

from .domain import SprintPlan


@dataclass(frozen=True)
class RemainingCapacitySamples:
    """Samples of remaining capacity S_t for one scenario."""

    samples: np.ndarray

    def __post_init__(self) -> None:
        s = np.asarray(self.samples, dtype=float).reshape(-1)
        object.__setattr__(self, "samples", s)

    @property
    def mean(self) -> float:
        return float(np.mean(self.samples)) if self.samples.size else 0.0

    @property
    def std(self) -> float:
        return float(np.std(self.samples, ddof=0)) if self.samples.size else 0.0


class CapacitySimulator:
    """Simulate remaining sprint capacity using the article's model.

    The simulator is intentionally focused on process uncertainty given a fixed
    set of parameters. Bayesian parameter uncertainty can be handled by an
    outer loop in the forecaster.
    """

    def __init__(self, plan: SprintPlan, *, beta_ai_focus: float):
        self.plan = plan
        self.beta_ai_focus = float(beta_ai_focus)

    def simulate_remaining_capacity(
        self,
        *,
        completed_days: int,
        n_samples: int,
        rng: np.random.Generator,
        override_dev_ai_intensity: Optional[float] = None,
    ) -> RemainingCapacitySamples:
        """Return samples of remaining capacity S_t.

        Parameters
        ----------
        completed_days:
            Number of completed days t. Remaining days are t..T-1.
        override_dev_ai_intensity:
            If provided, replaces all developer-day AI intensities u_{i,t}.
            Use 0.0 for the "noAI" counterfactual on the capacity side.
        """
        samples = self._simulate_pair(
            completed_days=completed_days,
            n_samples=n_samples,
            rng=rng,
            override_ai=override_dev_ai_intensity,
            override_noai=override_dev_ai_intensity,
        )[0]
        return RemainingCapacitySamples(samples=samples)

    def simulate_remaining_capacity_ai_and_noai(
        self,
        *,
        completed_days: int,
        n_samples: int,
        rng: np.random.Generator,
        override_ai: Optional[float] = None,
        override_noai: Optional[float] = 0.0,
    ) -> Tuple[RemainingCapacitySamples, RemainingCapacitySamples]:
        """Return (AI, noAI) samples using common random numbers."""
        s_ai, s0 = self._simulate_pair(
            completed_days=completed_days,
            n_samples=n_samples,
            rng=rng,
            override_ai=override_ai,
            override_noai=override_noai,
        )
        return RemainingCapacitySamples(samples=s_ai), RemainingCapacitySamples(samples=s0)

    def _simulate_pair(
        self,
        *,
        completed_days: int,
        n_samples: int,
        rng: np.random.Generator,
        override_ai: Optional[float],
        override_noai: Optional[float],
    ) -> Tuple[np.ndarray, np.ndarray]:
        plan = self.plan
        T = plan.n_days
        if not (0 <= completed_days <= T):
            raise ValueError(f"completed_days must be in [0,{T}]")
        if n_samples <= 0:
            raise ValueError("n_samples must be positive")
        if completed_days == T:
            zeros = np.zeros(n_samples, dtype=float)
            return zeros, zeros

        rem_days = np.arange(completed_days, T)
        h = plan.planned_hours[:, rem_days]  # (n, D)
        q = plan.availability_prob[:, rem_days]
        mu = plan.focus_log_mu[:, rem_days]
        sig = plan.focus_log_sigma[:, rem_days]
        u = plan.dev_day_ai_intensity[:, rem_days]

        if override_ai is not None:
            u_ai = np.full_like(u, float(override_ai))
        else:
            u_ai = u

        if override_noai is not None:
            u_0 = np.full_like(u, float(override_noai))
        else:
            u_0 = u

        # Sample blocked states per day.
        blocked_probs = plan.blocked.probabilities[:, rem_days]  # (K, D)
        g = plan.blocked.multipliers  # (K,)
        K = plan.blocked.n_states
        D = rem_days.size

        b_idx = np.empty((n_samples, D), dtype=int)
        for d in range(D):
            p = blocked_probs[:, d]
            b_idx[:, d] = rng.choice(K, size=n_samples, p=p)
        g_samples = g[b_idx]  # (n_samples, D)

        # Availability draws A_{i,day}.
        A = (rng.random((n_samples, plan.n_devs, D)) < q[None, :, :]).astype(float)

        # Common random numbers for focus (standard normal shocks).
        eps = rng.normal(loc=0.0, scale=1.0, size=(n_samples, plan.n_devs, D))

        # AI scenario focus.
        logF_ai = mu[None, :, :] + self.beta_ai_focus * u_ai[None, :, :] + sig[None, :, :] * eps
        F_ai = np.exp(logF_ai)

        # noAI scenario focus.
        logF_0 = mu[None, :, :] + self.beta_ai_focus * u_0[None, :, :] + sig[None, :, :] * eps
        F_0 = np.exp(logF_0)

        base_ai = np.sum(h[None, :, :] * A * F_ai, axis=1)
        base_0 = np.sum(h[None, :, :] * A * F_0, axis=1)

        daily_ai = base_ai * g_samples
        daily_0 = base_0 * g_samples

        s_ai = np.sum(daily_ai, axis=1)
        s_0 = np.sum(daily_0, axis=1)
        return s_ai, s_0
