"""Small demo of the sprint forecast model.

Run:
    python -m examples.demo

The numbers are illustrative only.
"""

from __future__ import annotations

import numpy as np

from sprint_forecast import BlockedStateSpec, BiasModelParameters, SprintPlan
from sprint_forecast import SprintSuccessForecaster, WorkItemEstimate


def main() -> None:
    rng = np.random.default_rng(42)

    developer_ids = ["alice", "bob", "carol", "dan"]
    n = len(developer_ids)
    T = 10  # working days

    # Planned hours: 4 hours/day per dev.
    h = np.full((n, T), 4.0)

    # Availability: mostly present.
    q = np.full((n, T), 0.95)

    # AI usage intensity for capacity channel (per dev-day).
    u_dev_day = rng.uniform(0.0, 1.0, size=(n, T))

    # Focus base parameters: lognormal on the log scale.
    focus_mu = np.full((n, T), np.log(0.65))
    focus_sigma = np.full((n, T), 0.25)

    # Blocked states: K=2 (normal day, incident day).
    g = np.array([1.0, 0.6])
    pi = np.vstack(
        [
            np.full(T, 0.85),  # normal
            np.full(T, 0.15),  # incident
        ]
    )
    blocked = BlockedStateSpec(multipliers=g, probabilities=pi)

    plan = SprintPlan(
        developer_ids=developer_ids,
        planned_hours=h,
        availability_prob=q,
        dev_day_ai_intensity=u_dev_day,
        focus_log_mu=focus_mu,
        focus_log_sigma=focus_sigma,
        blocked=blocked,
    )

    # Demand-side bias parameters (synthetic).
    bias = BiasModelParameters(
        mu_alpha=0.10,  # global optimism
        beta_ai=0.25,  # AI reduces actual effort relative to estimate
        alpha_by_engineer={"alice": 0.05, "bob": -0.02, "carol": 0.00, "dan": 0.03},
        delta_b=0.10,
    )

    # Remaining work items (synthetic) with an AI-assisted label/intensity.
    items = [
        WorkItemEstimate("A", "alice", 20.0, 28.0, 45.0, ai_intensity=1.0),
        WorkItemEstimate("B", "bob", 15.0, 22.0, 38.0, ai_intensity=0.0),
        WorkItemEstimate("C", "carol", 10.0, 18.0, 30.0, ai_intensity=0.6),
        WorkItemEstimate("D", "dan", 12.0, 20.0, 35.0, ai_intensity=0.8),
    ]

    forecaster = SprintSuccessForecaster(
        plan,
        bias=bias,
        beta_ai_focus=0.15,
    )

    completed_days = 5
    observed = 0.0  # using REMAINING_ONLY demand mode in this demo

    result = forecaster.estimate_ai_uplift(
        completed_days=completed_days,
        observed_cum_capacity=observed,
        remaining_items=items,
        n_samples=80_000,
        seed=123,
    )

    print(f"Completed days: {completed_days}/{T}")
    print(f"P(success | AI)            = {result.p_ai:.3f}")
    print(f"P(success | noAI)          = {result.p_no_ai:.3f}")
    print(f"Delta p (overall)          = {result.delta_p_overall:+.3f}")
    print(f"Delta p (workload channel) = {result.delta_p_from_workload:+.3f}")
    print(f"Delta p (capacity channel) = {result.delta_p_from_capacity:+.3f}")


if __name__ == "__main__":
    main()
