from __future__ import annotations

import random
from pathlib import Path

import torch

from deepgenopix.pixelizer import DNAPixelizer
from deepgenopix.quant.model import QuantEncoder
from deepgenopix.variant.null import (
    NullStats,
    benjamini_hochberg,
    encoder_fingerprint,
    reference_split_null,
    walk_redundancy_null,
)
from deepgenopix.variant.residual import call_variant_tokens, per_token_residual


def _tiny_encoder(seed: int = 0) -> QuantEncoder:
    torch.manual_seed(seed)
    encoder = QuantEncoder(
        num_classes=None,
        stride=2,
        latent_dim=32,
        n_heads=2,
        ffn_dim=64,
        n_layers=1,
        stem_channels=8,
        stem_kernel=3,
    )
    encoder.eval()
    return encoder


def _det_seq(length: int, seed: int) -> str:
    rng = random.Random(seed)
    return "".join(rng.choice("ACGT") for _ in range(length))


def test_benjamini_hochberg_thresholds_p_values():
    p = torch.tensor([0.001, 0.02, 0.5, 0.9])
    mask = benjamini_hochberg(p, q_target=0.05)
    assert mask.tolist() == [True, True, False, False]


def test_benjamini_hochberg_returns_empty_when_no_signal():
    p = torch.tensor([0.5, 0.6, 0.7])
    mask = benjamini_hochberg(p, q_target=0.05)
    assert not bool(mask.any())


def test_walk_redundancy_null_sigma_shrinks_with_stride():
    encoder = _tiny_encoder()
    pixelizer = DNAPixelizer(pixel_stride_bp=12)
    walk_pixelizer = DNAPixelizer(pixel_stride_bp=6)
    references = [_det_seq(144, seed=11), _det_seq(144, seed=12)]

    stats_nonwalk = walk_redundancy_null(encoder, references, pixelizer=pixelizer)
    stats_walk = walk_redundancy_null(encoder, references, pixelizer=walk_pixelizer)

    # Walk encoding doubles per-base coverage so σ should shrink ≈ 1/√2.
    assert stats_walk.sigma.mean().item() < stats_nonwalk.sigma.mean().item()
    assert stats_nonwalk.mu.abs().sum().item() == 0.0


def test_reference_split_null_runs_on_tiny_encoder():
    encoder = _tiny_encoder()
    pixelizer = DNAPixelizer()
    references = [_det_seq(300, seed=21), _det_seq(300, seed=22)]

    stats = reference_split_null(
        encoder,
        references,
        pixelizer=pixelizer,
        coverage=5,
        read_length=72,
        error_rate=0.0,
        n_samples=2,
        seed=1,
    )
    assert stats.mu.shape[0] == 2
    assert stats.sigma.shape == stats.mu.shape
    assert stats.estimator == "reference_split"
    assert (stats.sigma > 0).all()
    assert stats.encoder_hash == encoder_fingerprint(encoder)


def test_null_stats_round_trip(tmp_path: Path):
    encoder = _tiny_encoder()
    pixelizer = DNAPixelizer()
    references = [_det_seq(48, seed=31)]
    stats = walk_redundancy_null(encoder, references, pixelizer=pixelizer)
    path = tmp_path / "null.pt"
    stats.save(path)
    reloaded = NullStats.load(path)
    assert torch.allclose(reloaded.mu, stats.mu)
    assert torch.allclose(reloaded.sigma, stats.sigma)
    assert reloaded.encoder_hash == stats.encoder_hash
    assert reloaded.estimator == stats.estimator


def test_call_variant_tokens_with_null_stats_passes_only_outliers():
    n_tokens = 8
    mu = torch.zeros(1, n_tokens)
    sigma = torch.full((1, n_tokens), 0.05)
    stats = NullStats(
        mu=mu, sigma=sigma, encoder_hash="dummy", estimator="manual",
        coverage_per_locus=torch.tensor([0.0]),
    )
    # Build trajectories so that token 3 has a residual far above the null
    # and the rest sit at ~0.
    read = torch.zeros(n_tokens, 4)
    ref = torch.zeros(n_tokens, 4)
    read[:, 0] = 1.0
    ref[:, 0] = 1.0
    # Inject a clear mismatch only at token 3.
    read[3] = torch.tensor([0.0, 1.0, 0.0, 0.0])
    ref[3] = torch.tensor([1.0, 0.0, 0.0, 0.0])
    delta = per_token_residual(read, ref)
    assert float(delta[3].item()) > 0.5

    calls = call_variant_tokens(
        read, ref, locus_idx=0, null_stats=stats, q_target=0.05,
    )
    assert 3 in calls
    # No other tokens should pass (they sit at Δ ≈ 0).
    assert calls == [3]


def test_call_variant_tokens_legacy_path_still_works():
    n_tokens = 4
    read = torch.zeros(n_tokens, 4)
    ref = torch.zeros(n_tokens, 4)
    read[:, 0] = 1.0
    ref[:, 0] = 1.0
    read[2] = torch.tensor([0.0, 1.0, 0.0, 0.0])
    ref[2] = torch.tensor([1.0, 0.0, 0.0, 0.0])
    calls = call_variant_tokens(read, ref, threshold=0.5)
    assert calls == [2]