import numpy as np
import pytest
import torch

import _net
import ssl_mamba as S


def test_pad_long_truncated():
    x = np.arange(100000, dtype=np.float32)
    assert np.array_equal(_net.pad_repeat(x), x[:64600])


def test_pad_short_repeated():
    x = np.arange(1000, dtype=np.float32)
    out = _net.pad_repeat(x)
    assert out.shape == (64600,) and np.array_equal(out[:1000], x) and np.array_equal(out[1000:2000], x)


@pytest.fixture(scope="module")
def model():
    m = S.SSLMamba(); m.load()   # strict=True load (0 missing / unexpected)
    yield m
    m.unload()


def test_eval_mode_and_params(model):
    assert not any(mod.training for mod in model.net.modules())
    assert sum(p.numel() for p in model.net.parameters()) == 317_521_282
    assert sum(t.numel() for t in model.net.state_dict().values()) == 317_521_282  # no buffers


def test_batch_matches_single(model):
    rng = np.random.default_rng(0)
    a = [rng.standard_normal(n).astype(np.float32) * s for n, s in ((48000, 0.1), (20000, 0.5), (70000, 0.02), (3000, 0.3))]
    batched = model.score_batch(a, [16000] * 4)
    single = [model.score_batch([x], [16000])[0] for x in a]
    assert len(batched) == 4 and all(np.isfinite(batched))
    # cuDNN picks batch-dependent conv algorithms (TF32 on by default) -> small relative jitter
    assert np.allclose(batched, single, rtol=1e-4, atol=1e-4)


def test_chunked_scan_matches_reference(model):
    rng = np.random.default_rng(1)
    a = [rng.standard_normal(50000).astype(np.float32) * 0.1 for _ in range(3)]
    s1 = model.score_batch(a, [16000] * 3)
    model.net.set_ref_scan(True)
    try:
        s2 = model.score_batch(a, [16000] * 3)
    finally:
        model.net.set_ref_scan(False)
    assert np.allclose(s1, s2, atol=1e-4)


def test_deterministic(model):
    a = [np.random.randn(40000).astype(np.float32) for _ in range(2)]
    assert model.score_batch(a, [16000] * 2) == model.score_batch(a, [16000] * 2)


def test_sign_and_finite(model):
    # higher = bona fide: score equals log-odds of class 1 (bona fide) vs class 0
    a = [np.random.default_rng(2).standard_normal(64600).astype(np.float32) * 0.05]
    s = model.score_batch(a, [16000])[0]
    x = torch.from_numpy(_net.pad_repeat(a[0])[None]).to(model.device)
    with torch.no_grad():
        out = model.net(x)[0]
    assert np.isfinite(s) and abs(s - float(out[1] - out[0])) < 1e-4