# tests/test_screen.py — Tests for the fast screening module.
"""
Tests for geney.screen — fast variant screening with minimal-computation deltas.
"""
import pytest
import numpy as np


# ============================================================================
# Helper
# ============================================================================

def _seqmat_available() -> bool:
    """Check if seqmat is installed and data is available."""
    try:
        from seqmat import Gene
        Gene.from_file("KRAS", organism="hg38")
        return True
    except Exception:
        return False


_HAS_SEQMAT = _seqmat_available()


# ============================================================================
# Unit tests for individual screen functions
# ============================================================================

class TestScreenResult:
    def test_max_delta_all_none(self):
        from geney.screen import ScreenResult
        r = ScreenResult("X:1:100:A:T", "X", None, None, None, None)
        assert r.max_delta == 0.0

    def test_max_delta_mixed(self):
        from geney.screen import ScreenResult
        r = ScreenResult("X:1:100:A:T", "X", 0.3, -0.5, 0.1, None)
        assert r.max_delta == 0.5

    def test_max_delta_all_set(self):
        from geney.screen import ScreenResult
        r = ScreenResult("X:1:100:A:T", "X", 0.1, 0.2, 0.3, 0.9)
        assert r.max_delta == 0.9


class TestScreenSplicingFromProbs:
    def test_identical_probs(self):
        from geney.screen import screen_splicing_from_probs
        probs = np.array([0.1, 0.5, 0.9])
        result = screen_splicing_from_probs(probs, probs, probs, probs)
        assert result == 0.0

    def test_different_probs(self):
        from geney.screen import screen_splicing_from_probs
        ref_don = np.array([0.1, 0.5, 0.9])
        var_don = np.array([0.1, 0.8, 0.9])
        ref_acc = np.zeros(3)
        var_acc = np.zeros(3)
        result = screen_splicing_from_probs(ref_don, ref_acc, var_don, var_acc)
        assert abs(result - 0.3) < 1e-6


class TestScreenFolding:
    def test_identical_sequences(self):
        from geney.screen import screen_folding
        seq = "ACGUACGUACGUACGUACGUACGUACGUACGUACGUACGUA"
        result = screen_folding(seq, seq, 20, window_size=39)
        assert result == 0.0

    def test_out_of_bounds(self):
        from geney.screen import screen_folding
        result = screen_folding("ACGU", "ACGU", -1)
        assert result == 0.0
        result = screen_folding("ACGU", "ACGU", 10)
        assert result == 0.0


class TestScreenTis:
    def test_no_atg_returns_zero(self):
        from geney.screen import screen_tis
        result = screen_tis("CCCCCCCCCCC", "CCCCCCCCCCC")
        assert result == 0.0


# ============================================================================
# Integration tests (require seqmat data)
# ============================================================================

@pytest.mark.skipif(not _HAS_SEQMAT, reason="seqmat data not available")
def test_screen_mutation_kras():
    """Smoke test: screen a KRAS mutation."""
    from geney.screen import screen_mutation
    result = screen_mutation(
        "KRAS:12:25227343:G:T",
        run_splicing=True,
        run_tis=True,
        run_folding=True,
        run_transcription=False,
    )
    assert result.mut_id == "KRAS:12:25227343:G:T"
    assert result.gene == "KRAS"
    assert result.error is None
    assert result.splicing_delta is not None
    assert result.splicing_delta >= 0.0


@pytest.mark.skipif(not _HAS_SEQMAT, reason="seqmat data not available")
def test_screen_mutations_streaming():
    """Test streaming screen with multiple mutations."""
    from geney.screen import screen_mutations
    mut_ids = [
        "KRAS:12:25227343:G:T",
        "KRAS:12:25227342:C:A",
    ]
    results = list(screen_mutations(
        mut_ids,
        batch_size=2,
        run_splicing=True,
        run_tis=False,
        run_folding=False,
        run_transcription=False,
    ))
    assert len(results) == 2
    for r in results:
        assert r.gene == "KRAS"


@pytest.mark.skipif(not _HAS_SEQMAT, reason="seqmat data not available")
def test_screen_and_select():
    """Test top-K selection."""
    from geney.screen import screen_and_select
    mut_ids = [
        "KRAS:12:25227343:G:T",
        "KRAS:12:25227342:C:A",
        "KRAS:12:25227341:T:G",
    ]
    top = screen_and_select(
        mut_ids, top_k=2,
        run_splicing=True,
        run_tis=False,
        run_folding=False,
        run_transcription=False,
    )
    assert len(top) <= 2
    if len(top) == 2:
        assert top[0].max_delta >= top[1].max_delta


@pytest.mark.skipif(not _HAS_SEQMAT, reason="seqmat data not available")
def test_on_result_callback():
    """Test that on_result callback fires for each mutation."""
    from geney.screen import screen_mutations
    collected = []

    list(screen_mutations(
        ["KRAS:12:25227343:G:T"],
        run_splicing=True,
        run_tis=False,
        run_folding=False,
        run_transcription=False,
        on_result=collected.append,
    ))

    assert len(collected) == 1
    assert collected[0].gene == "KRAS"
