import numpy as np
import torch

from chronos_x.models.joint_path import samples_from_copula, LowRankCopulaHead, pit_student_t, z_from_pit


def test_copula_samples_reproduce_marginals():
    torch.manual_seed(0)
    n, B, H = 20000, 2, 60
    z = torch.randn(n, B, H)
    mu = torch.full((B, H), 0.001); sc = torch.full((B, H), 0.002); nu = torch.full((B, H), 6.0)
    S = samples_from_copula(z, mu, sc, nu)
    assert abs(float(S.mean()) - 0.001) < 5e-5
    assert abs(float(S.std()) - 0.002 * np.sqrt(6 / 4)) < 1e-4


def test_sample_z_is_standardised_and_correlation_consistent():
    torch.manual_seed(1)
    head = LowRankCopulaHead(8, 60, rank=2)
    with torch.no_grad():
        head.load.weight.mul_(20.0)        # strong common factor
    h = torch.randn(3, 60, 8)
    z = head.sample_z(h, 20000)
    R = head.correlation(h)
    emp = torch.corrcoef(z[:, 0, :].T)
    assert abs(float(z.std()) - 1) < 0.02
    assert float((emp - R[0]).abs().mean()) < 0.03


def test_copula_nll_minimised_at_independence_for_independent_data():
    torch.manual_seed(2)
    head = LowRankCopulaHead(8, 60, rank=2)
    h = torch.randn(64, 60, 8)
    z = torch.randn(64, 60); mask = torch.ones(64, 60, dtype=torch.bool)
    nll_small = head.nll(h, z, mask)
    with torch.no_grad():
        head.load.weight.mul_(30.0)
    nll_big = head.nll(h, z, mask)
    assert nll_small < nll_big


def test_pit_roundtrip():
    y = np.array([[-0.003, 0.0, 0.004]]); mu = np.zeros((1, 3)); sc = np.full((1, 3), 0.002); nu = np.full((1, 3), 5.0)
    u = pit_student_t(y, mu, sc, nu)
    assert 0 < u[0, 0] < 0.5 < u[0, 2] < 1 and abs(u[0, 1] - 0.5) < 1e-12
    z = z_from_pit(u)
    assert z[0, 0] < 0 < z[0, 2]