import numpy as np
import torch

from chronos_x.models.model import MRPT, MRPTConfig
from chronos_x.data.windows import F_BAR, F_GLOBAL
from chronos_x.models.student_t import masked_student_t_nll, student_t_logpdf


def _batch(B=3, H=60, LA=390, LB=390, LC=260):
    g = torch.Generator().manual_seed(0)
    return {"xA": torch.randn(B, LA, F_BAR, generator=g), "mA": torch.ones(B, LA, dtype=torch.uint8),
            "xB": torch.randn(B, LB, F_BAR, generator=g), "mB": torch.ones(B, LB, dtype=torch.uint8),
            "xC": torch.randn(B, LC, F_BAR, generator=g), "mC": torch.ones(B, LC, dtype=torch.uint8),
            "glob": torch.randn(B, F_GLOBAL, generator=g), "fut_cal": torch.rand(B, H, 4, generator=g),
            "sigma": torch.full((B,), 3e-4), "y": torch.randn(B, H, generator=g) * 3e-4,
            "y_mask": torch.ones(B, H, dtype=torch.uint8), "sym": torch.zeros(B, dtype=torch.long)}


def test_forward_shapes_and_param_counts():
    for preset, lo, hi in (("tiny", 1e4, 5e5), ("small", 2e6, 6e6), ("spec", 5e6, 12e6), ("base", 15e6, 30e6)):
        m = MRPT(MRPTConfig.preset(preset))
        n = m.n_params()
        assert lo <= n <= hi, (preset, n)
    m = MRPT(MRPTConfig.preset("tiny"))
    out = m(_batch())
    assert out["mu"].shape == (3, 60) and out["scale"].shape == (3, 60) and out["nu"].shape == (3, 60)
    assert (out["scale"] > 0).all() and (out["nu"] > 2).all()
    loss = masked_student_t_nll(_batch()["y"], _batch()["y_mask"], out)
    assert torch.isfinite(loss)
    loss.backward()


def test_variants_run():
    b = _batch()
    for kw in (dict(branches=["A"]), dict(branches=["A", "B"]), dict(fusion="cross"), dict(head="gaussian"),
               dict(head="quantile"), dict(regime="none"), dict(regime="learned"), dict(n_symbols=4),
               dict(copula_rank=3)):
        m = MRPT(MRPTConfig.preset("tiny", **kw))
        out = m(b)
        key = "quantiles" if kw.get("head") == "quantile" else "mu"
        assert torch.isfinite(out[key]).all()


def test_student_t_logpdf_matches_scipy():
    from scipy import stats
    y = torch.tensor([[-0.002, 0.0, 0.0031]]); mu = torch.tensor([[0.0001, 0.0, -0.0002]])
    sc = torch.tensor([[0.001, 0.002, 0.0005]]); nu = torch.tensor([[3.0, 5.5, 20.0]])
    ours = student_t_logpdf(y, mu, sc, nu).numpy()
    ref = stats.t.logpdf(y.numpy(), nu.numpy(), loc=mu.numpy(), scale=sc.numpy())
    np.testing.assert_allclose(ours, ref, rtol=1e-5)