tests/test_variant_simulate.py

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

from deepgenopix.variant.simulate import (
    SimulatorConfig,
    VariantManifest,
    inject_variants,
    materialize_corpus,
    reverse_complement,
    simulate_reads_with_truth,
    token_position_for,
)


def _det_seq(length: int, seed: int) -> str:
    rng = np.random.default_rng(seed)
    return "".join("ACGT"[int(b)] for b in rng.integers(0, 4, size=length))


def test_inject_snp_changes_only_target_position():
    seq = _det_seq(120, seed=11)
    pos = 37
    ref_base = seq[pos]
    alt_base = "ACGT"[("ACGT".index(ref_base) + 1) % 4]
    manifest = VariantManifest("locus_a", pos, "snp", ref_base, alt_base, 0.5)
    out = inject_variants(seq, [manifest])
    assert out[pos] == alt_base
    assert out[:pos] == seq[:pos]
    assert out[pos + 1 :] == seq[pos + 1 :]


def test_inject_insertion_grows_sequence():
    seq = _det_seq(96, seed=22)
    insert = "GGGAAACCC"
    manifest = VariantManifest("locus_a", 40, "ins", "", insert, 0.5)
    out = inject_variants(seq, [manifest])
    assert len(out) == len(seq) + len(insert)
    assert out[40 : 40 + len(insert)] == insert
    assert out[: 40] == seq[: 40]
    assert out[40 + len(insert) :] == seq[40:]


def test_inject_deletion_removes_window():
    seq = _det_seq(96, seed=33)
    start = 50
    deletion = seq[start : start + 4]
    manifest = VariantManifest("locus_a", start, "del", deletion, "", 0.5)
    out = inject_variants(seq, [manifest])
    assert len(out) == len(seq) - 4
    assert out[: start] == seq[: start]
    assert out[start :] == seq[start + 4 :]


def test_inject_inversion_reverse_complements_window():
    seq = _det_seq(96, seed=44)
    start = 24
    length = 12
    segment = seq[start : start + length]
    manifest = VariantManifest("locus_a", start, "inv", segment, reverse_complement(segment), 0.5)
    out = inject_variants(seq, [manifest])
    assert out[start : start + length] == reverse_complement(segment)
    assert out[: start] == seq[: start]
    assert out[start + length :] == seq[start + length :]


def test_simulate_reads_truth_round_trip_for_snp():
    seq = _det_seq(300, seed=55)
    pos = 64
    ref_base = seq[pos]
    alt_base = "ACGT"[("ACGT".index(ref_base) + 1) % 4]
    manifest = VariantManifest("locus_a", pos, "snp", ref_base, alt_base, 0.5)

    reads, truth = simulate_reads_with_truth(
        seq,
        [manifest],
        coverage=40,
        read_length=120,
        error_rate=0.0,
        seed=7,
        haplotype_frequency=0.5,
        bp_per_token=24,
    )
    assert len(reads) == len(truth)
    assert set(truth["haplotype"]) == {"variant", "reference"}

    # Reads that cover the variant position from the variant haplotype must
    # show the alt base; reference reads must show the ref base.
    variant_rows = truth[truth["haplotype"] == "variant"]
    for _, row in variant_rows.iterrows():
        start = int(row["read_offset_in_haplotype"])
        end = start + 120
        if not (start <= pos < end):
            continue
        offset_in_read = pos - start
        assert reads[int(row["read_idx"])][offset_in_read] == alt_base
        assert row["manifest_idxs"] == (0,)
        assert row["variant_token_positions"] == (token_position_for(pos, bp_per_token=24),)

    reference_rows = truth[truth["haplotype"] == "reference"]
    for _, row in reference_rows.iterrows():
        start = int(row["read_offset_in_haplotype"])
        end = start + 120
        if not (start <= pos < end):
            continue
        offset_in_read = pos - start
        assert reads[int(row["read_idx"])][offset_in_read] == ref_base


def test_token_position_for_respects_bp_per_token():
    assert token_position_for(0, bp_per_token=24) == 0
    assert token_position_for(23, bp_per_token=24) == 0
    assert token_position_for(24, bp_per_token=24) == 1
    assert token_position_for(60, bp_per_token=24, locus_offset=12) == 2


def test_materialize_corpus_round_trips(tmp_path: Path):
    sequences = [_det_seq(300, seed=100 + i) for i in range(3)]
    reference_parquet = tmp_path / "reference.parquet"
    pd.DataFrame(
        {
            "locus_id": [f"locus_{i}" for i in range(len(sequences))],
            "sequence": sequences,
        }
    ).to_parquet(reference_parquet, index=False)

    output_dir = tmp_path / "corpus"
    cfg = SimulatorConfig(
        coverage=20,
        read_length=150,
        error_rate=0.0,
        bp_per_token=24,
        seed=42,
        variants_per_locus=2,
        variant_types=("snp",),
        frequency_levels=(0.5,),
    )
    paths = materialize_corpus(reference_parquet, cfg, output_dir)
    assert paths.reads.exists()
    assert paths.truth.exists()
    assert paths.manifest.exists()
    assert paths.summary.exists()

    reads_df = pd.read_parquet(paths.reads)
    truth_df = pd.read_parquet(paths.truth)
    manifests_df = pd.read_parquet(paths.manifest)
    assert len(reads_df) == len(truth_df)
    assert len(manifests_df) == cfg.variants_per_locus * len(sequences)
    assert truth_df["haplotype"].isin({"reference", "variant"}).all()


def test_materialize_corpus_skips_loci_that_become_too_short(tmp_path: Path):
    reference_parquet = tmp_path / "reference.parquet"
    pd.DataFrame(
        {
            "locus_id": ["too_short_after_deletion", "long_enough"],
            "sequence": [_det_seq(151, seed=11), _det_seq(220, seed=12)],
        }
    ).to_parquet(reference_parquet, index=False)

    cfg = SimulatorConfig(
        coverage=5,
        read_length=150,
        error_rate=0.0,
        seed=0,
        variant_types=("del",),
        indel_size_bp=3,
    )
    paths = materialize_corpus(reference_parquet, cfg, tmp_path / "corpus")
    summary = json.loads(paths.summary.read_text())

    assert summary["n_loci"] == 1
    assert summary["n_skipped_loci"] == 1
    assert "haplotype shorter than read_length" in summary["skipped_loci"][0]["reason"]