import numpy as np
import pandas as pd
from scipy.stats import norm

from btcmodel.forecast import Forecast, crps, pit


def normal_crps(mu, sigma, y):
    """Closed-form CRPS of N(mu, sigma^2) at y."""
    z = (y - mu) / sigma
    return sigma * (z * (2 * norm.cdf(z) - 1) + 2 * norm.pdf(z) - 1 / np.sqrt(np.pi))


def test_crps_matches_closed_form_for_normal():
    f = Forecast.normal("2020-01-01", np.array([1, 2, 3]), mean=[0.0, 1.0, 2.0], sd=[1.0, 0.5, 2.0])
    y = np.array([0.3, -0.2, 5.0])
    expected = normal_crps(np.array([0.0, 1.0, 2.0]), np.array([1.0, 0.5, 2.0]), y)
    np.testing.assert_allclose(crps(f.log_quantiles, y), expected, rtol=0.02)


def test_crps_prefers_the_right_forecast():
    y = np.array([0.0])
    good = Forecast.normal("2020-01-01", np.array([1]), 0.0, 0.1)
    biased = Forecast.normal("2020-01-01", np.array([1]), 0.5, 0.1)
    vague = Forecast.normal("2020-01-01", np.array([1]), 0.0, 2.0)
    assert crps(good.log_quantiles, y) < crps(biased.log_quantiles, y)
    assert crps(good.log_quantiles, y) < crps(vague.log_quantiles, y)


def test_pit_and_intervals():
    f = Forecast.normal("2020-01-01", np.array([1, 1, 1]), 0.0, 1.0)
    np.testing.assert_allclose(
        pit(f.log_quantiles, [0.0, 1.0, -10.0]), [0.5, 0.841, 0.0], atol=0.01
    )
    lo, hi = f.interval(0.95)
    np.testing.assert_allclose(hi, 1.96, atol=0.01)
    np.testing.assert_allclose(lo, -1.96, atol=0.01)


def test_from_samples_recovers_quantiles():
    rng = np.random.default_rng(0)
    f = Forecast.from_samples("2020-01-01", np.array([10]), rng.normal(0, 1, size=(200_000, 1)))
    np.testing.assert_allclose(f.quantile(0.5), 0.0, atol=0.02)
    assert f.dates[0] == pd.Timestamp("2020-01-11")