import unittest

import numpy as np

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


class TestForecaster(unittest.TestCase):
    def test_ai_uplift_is_positive_in_borderline_case(self) -> None:
        developer_ids = ["a", "b"]
        n = len(developer_ids)
        T = 5

        # Low planned hours so success is non-trivial without AI.
        h = np.full((n, T), 3.0)
        q = np.full((n, T), 1.0)  # always available

        # AI usage positive on all days.
        u_dev_day = np.ones((n, T))

        focus_mu = np.full((n, T), np.log(0.7))
        focus_sigma = np.full((n, T), 0.15)

        # Single blocked state.
        g = np.array([1.0])
        pi = np.ones((1, T))
        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,
        )

        bias = BiasModelParameters(
            mu_alpha=0.0,
            beta_ai=0.5,
            alpha_by_engineer={"a": 0.0, "b": 0.0},
            delta_b=0.0,
        )

        # A single big item that is borderline without AI.
        items = [
            WorkItemEstimate("X", "a", 18.0, 22.0, 30.0, ai_intensity=1.0),
        ]

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

        result = forecaster.estimate_ai_uplift(
            completed_days=0,
            observed_cum_capacity=0.0,
            remaining_items=items,
            n_samples=20000,
            seed=123,
        )

        self.assertGreater(result.delta_p_overall, 0.05)


if __name__ == "__main__":
    unittest.main()
