import unittest

import numpy as np

from sprint_forecast.models.bias import (
    BiasHyperparameters,
    BiasObservation,
    HierarchicalBiasEstimator,
)


class TestHierarchicalBiasEstimator(unittest.TestCase):
    def test_estimates_ai_coefficient_reasonably(self) -> None:
        rng = np.random.default_rng(1)
        engineer_ids = ["a", "b"]

        mu_alpha = 0.2
        alpha = {"a": 0.1, "b": -0.05}
        beta_ai = 0.3
        sigma = 0.05

        obs = []
        for eid in engineer_ids:
            for u in [0.0, 1.0]:
                for _ in range(50):
                    eps = rng.normal(0.0, sigma)
                    eta = mu_alpha + alpha[eid] - beta_ai * u + eps
                    # choose estimate_m=1, actual=exp(eta)
                    obs.append(BiasObservation(eid, 1.0, float(np.exp(eta)), ai_intensity=u))

        hyper = BiasHyperparameters(
            mu_prior_mean=0.0,
            mu_prior_std=1.0,
            beta_ai_prior_mean=0.0,
            beta_ai_prior_std=1.0,
            alpha_prior_std=1.0,
            noise_std=sigma,
        )
        est = HierarchicalBiasEstimator(hyper)
        params = est.fit(obs, engineer_ids)

        self.assertAlmostEqual(params.beta_ai, beta_ai, delta=0.05)


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