import unittest

import numpy as np

from sprint_forecast.distributions import QuantileLevels, WeightedLogNormalQuantileFitter


class TestLogNormalQuantileFitter(unittest.TestCase):
    def test_recovers_parameters_from_quantiles(self) -> None:
        mu_true = 1.2
        sigma_true = 0.4
        levels = QuantileLevels(0.1, 0.5, 0.9)
        fitter = WeightedLogNormalQuantileFitter(levels=levels)

        # Generate exact quantiles for the underlying lognormal.
        from statistics import NormalDist

        nd = NormalDist(0.0, 1.0)
        z = np.array([nd.inv_cdf(p) for p in levels.as_tuple])
        q = np.exp(mu_true + sigma_true * z)
        a, m, b = float(q[0]), float(q[1]), float(q[2])

        params = fitter.fit(a=a, m=m, b=b)

        self.assertAlmostEqual(params.mu, mu_true, places=6)
        self.assertAlmostEqual(params.sigma, sigma_true, places=6)


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