"""Top-level API for probability of sprint success and AI attribution."""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Sequence

import numpy as np

from .capacity import CapacitySimulator
from .domain import SprintPlan, WorkItemEstimate
from .models.bias import BiasModelParameters
from .workload import WorkloadModelBuilder


class DemandMode(str, Enum):
    """How to interpret workload samples W when forming residual demand."""

    REMAINING_ONLY = "remaining_only"
    TOTAL_MINUS_DELIVERED = "total_minus_delivered"


@dataclass(frozen=True)
class AICounterfactualResult:
    """Estimated success probabilities under AI and counterfactual settings."""

    completed_days: int
    n_samples: int

    p_ai: float
    p_no_ai: float

    p_no_ai_workload: float
    p_no_ai_capacity: float

    @property
    def delta_p_overall(self) -> float:
        return self.p_ai - self.p_no_ai

    @property
    def delta_p_from_workload(self) -> float:
        """Approximate marginal uplift from AI in the workload channel."""
        return self.p_ai - self.p_no_ai_workload

    @property
    def delta_p_from_capacity(self) -> float:
        """Approximate marginal uplift from AI in the capacity channel."""
        return self.p_ai - self.p_no_ai_capacity


class SprintSuccessForecaster:
    """Compute the probability of sprint success on day t via paired simulation."""

    def __init__(
        self,
        plan: SprintPlan,
        *,
        bias: BiasModelParameters,
        beta_ai_focus: float,
        workload_builder: WorkloadModelBuilder | None = None,
    ):
        self.plan = plan
        self.bias = bias
        self.capacity_sim = CapacitySimulator(plan, beta_ai_focus=beta_ai_focus)
        self.workload_builder = workload_builder or WorkloadModelBuilder()

    def estimate_success_probability(
        self,
        *,
        completed_days: int,
        observed_cum_capacity: float,
        remaining_items: Sequence[WorkItemEstimate],
        n_samples: int = 100_000,
        seed: int | None = None,
        demand_mode: DemandMode = DemandMode.REMAINING_ONLY,
    ) -> float:
        """Estimate p_succ(t) = P(S_t >= R_t).

        Parameters
        ----------
        completed_days:
            Number of completed days t.
        observed_cum_capacity:
            Cumulative realised capacity \\hat{C}_t.
        remaining_items:
            Items to include in the workload model.
        demand_mode:
            - REMAINING_ONLY: treat workload samples as already residual; R_t = W.
            - TOTAL_MINUS_DELIVERED: treat workload samples as total sprint demand;
              R_t = max(0, W - \\hat{C}_t).
        """
        rng = np.random.default_rng(seed)

        workload_sampler = self.workload_builder.build_sampler(
            remaining_items, self.bias
        )
        w = workload_sampler.sample_total(rng=rng, size=n_samples)

        s = self.capacity_sim.simulate_remaining_capacity(
            completed_days=completed_days,
            n_samples=n_samples,
            rng=rng,
        ).samples

        r = _residual_demand(w, observed_cum_capacity, mode=demand_mode)
        return float(np.mean(s >= r))

    def estimate_ai_uplift(
        self,
        *,
        completed_days: int,
        observed_cum_capacity: float,
        remaining_items: Sequence[WorkItemEstimate],
        n_samples: int = 100_000,
        seed: int | None = None,
        demand_mode: DemandMode = DemandMode.REMAINING_ONLY,
    ) -> AICounterfactualResult:
        """Estimate success probabilities with AI and counterfactual no-AI.

        This returns:
        - p_ai:    AI used in both workload + capacity channels.
        - p_no_ai: AI removed from both channels.
        - p_no_ai_workload: AI removed only from workload model.
        - p_no_ai_capacity: AI removed only from capacity model.

        All four are computed from the same random draws (common random numbers)
        to reduce Monte Carlo noise.
        """
        rng = np.random.default_rng(seed)

        # Workload samplers for AI and noAI demand-side assumptions.
        wl_ai = self.workload_builder.build_sampler(
            remaining_items, self.bias, override_item_ai_intensity=None
        )
        wl_0 = self.workload_builder.build_sampler(
            remaining_items, self.bias, override_item_ai_intensity=0.0
        )

        # Generate common random normals for all workload scenarios.
        if wl_ai.n_items != wl_0.n_items:
            raise RuntimeError("Internal error: workload samplers have different sizes")

        z_w = (
            rng.normal(loc=0.0, scale=1.0, size=(n_samples, wl_ai.n_items))
            if wl_ai.n_items
            else np.zeros((n_samples, 0), dtype=float)
        )

        w_ai = wl_ai.sample_total_with_normals(z_w)
        w_0 = wl_0.sample_total_with_normals(z_w)

        # Capacity samples for AI and noAI capacity channels.
        s_ai, s_0 = self.capacity_sim.simulate_remaining_capacity_ai_and_noai(
            completed_days=completed_days,
            n_samples=n_samples,
            rng=rng,
            override_ai=None,
            override_noai=0.0,
        )

        r_ai = _residual_demand(w_ai, observed_cum_capacity, mode=demand_mode)
        r_0 = _residual_demand(w_0, observed_cum_capacity, mode=demand_mode)

        # Crossed counterfactuals.
        # - noAI workload only: use W_0 with capacity S_AI.
        # - noAI capacity only: use W_AI with capacity S_0.
        r_no_ai_workload = _residual_demand(w_0, observed_cum_capacity, mode=demand_mode)
        r_no_ai_capacity = _residual_demand(w_ai, observed_cum_capacity, mode=demand_mode)

        p_ai = float(np.mean(s_ai.samples >= r_ai))
        p_no_ai = float(np.mean(s_0.samples >= r_0))
        p_no_ai_workload = float(np.mean(s_ai.samples >= r_no_ai_workload))
        p_no_ai_capacity = float(np.mean(s_0.samples >= r_no_ai_capacity))

        return AICounterfactualResult(
            completed_days=completed_days,
            n_samples=n_samples,
            p_ai=p_ai,
            p_no_ai=p_no_ai,
            p_no_ai_workload=p_no_ai_workload,
            p_no_ai_capacity=p_no_ai_capacity,
        )


def _residual_demand(
    w: np.ndarray,
    observed_cum_capacity: float,
    *,
    mode: DemandMode,
) -> np.ndarray:
    if mode == DemandMode.REMAINING_ONLY:
        return np.maximum(0.0, w)
    if mode == DemandMode.TOTAL_MINUS_DELIVERED:
        return np.maximum(0.0, w - float(observed_cum_capacity))
    raise ValueError(f"Unknown DemandMode: {mode}")
