tests/test_variant_decoder.py

from __future__ import annotations

import random

import pandas as pd

from deepgenopix.pixelizer import DNAPixelizer
from deepgenopix.variant import call_cna, pixel_to_12mer
from deepgenopix.variant.decoder import (
    base_at_intra_offset,
    base_position_to_pixel_offsets,
    offsets_for_strip,
    strip_to_sequence,
    vote_base_at_position,
)


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


def test_variant_decoder_round_trip_random_12mer():
    sequence = _det_seq(12, seed=1)
    pixelizer = DNAPixelizer()

    tensor, n_pixels = pixelizer.string_to_tensor(sequence)

    assert n_pixels == 1
    assert pixel_to_12mer(tensor[:, 0]) == sequence


def test_call_cna_returns_expected_calls():
    tumor = pd.Series({"a": 30, "b": 10, "c": 10})
    normal = pd.Series({"a": 10, "b": 10, "c": 30})

    calls = call_cna(tumor, normal)

    by_locus = dict(zip(calls["locus"], calls["call"], strict=False))
    assert by_locus["a"] == "amp"
    assert by_locus["c"] == "del"


def test_strip_to_sequence_recovers_at_stride_12():
    sequence = _det_seq(96, seed=2)
    pixelizer = DNAPixelizer()
    tensor, n_pixels = pixelizer.string_to_tensor(sequence)

    decoded, disagreement = strip_to_sequence(tensor, pixel_stride_bp=12, sequence_length=len(sequence))
    assert decoded == sequence
    # interior bases observed by exactly one pixel → no disagreement signal
    assert all(d == 0 for d in disagreement)
    assert n_pixels == len(sequence) // 12


def test_strip_to_sequence_recovers_at_walk_stride_9():
    sequence = _det_seq(99, seed=3)
    pixelizer = DNAPixelizer(pixel_stride_bp=9)
    tensor, _ = pixelizer.string_to_tensor(sequence)

    decoded, disagreement = strip_to_sequence(tensor, pixel_stride_bp=9, sequence_length=len(sequence))
    assert decoded == sequence
    # at stride 9 the walk overlaps so several positions get >1 vote; under a
    # clean encoding the votes all agree → disagreement counts stay at zero
    assert max(disagreement) == 0


def test_strip_to_sequence_recovers_at_stride_1():
    sequence = _det_seq(36, seed=4)
    pixelizer = DNAPixelizer(pixel_stride_bp=1)
    tensor, _ = pixelizer.string_to_tensor(sequence)

    decoded, disagreement = strip_to_sequence(tensor, pixel_stride_bp=1, sequence_length=len(sequence))
    assert decoded == sequence
    # interior bases at stride 1 are observed by 12 pixels; under clean encoding
    # all 12 votes agree, so disagreement is zero
    interior_disagreement = disagreement[12:-12]
    assert all(d == 0 for d in interior_disagreement)


def test_offsets_for_strip_matches_pixelizer():
    sequence = _det_seq(73, seed=5)
    pixelizer = DNAPixelizer(pixel_stride_bp=9)
    tensor, n_pixels = pixelizer.string_to_tensor(sequence)

    offsets = offsets_for_strip(n_pixels, pixel_stride_bp=9, sequence_length=len(sequence))
    assert len(offsets) == n_pixels
    # the last offset must always anchor at sequence_end - 12 (snap-back)
    assert offsets[-1] == len(sequence) - 12


def test_base_position_to_pixel_offsets_covers_every_base():
    sequence = _det_seq(45, seed=6)
    pixelizer = DNAPixelizer(pixel_stride_bp=6)
    _, n_pixels = pixelizer.string_to_tensor(sequence)
    offsets = offsets_for_strip(n_pixels, pixel_stride_bp=6, sequence_length=len(sequence))

    for base_position in range(len(sequence)):
        covers = base_position_to_pixel_offsets(offsets, base_position)
        assert covers, f"position {base_position} has no covering pixel"
        for pixel_idx, intra in covers:
            assert 0 <= intra < 12
            assert offsets[pixel_idx] + intra == base_position


def test_disagreement_signals_corrupted_pixel():
    sequence = _det_seq(36, seed=7)
    pixelizer = DNAPixelizer(pixel_stride_bp=6)
    tensor, n_pixels = pixelizer.string_to_tensor(sequence)
    offsets = offsets_for_strip(n_pixels, pixel_stride_bp=6, sequence_length=len(sequence))

    # Corrupt one pixel and confirm vote disagreement appears at every position
    # that pixel covered.
    target_pixel = 1
    original_base = base_at_intra_offset(tensor[:, target_pixel], 0)
    flipped = "ACGT"[("ACGT".index(original_base) + 1) % 4]
    # Build a new packed value whose intra-offset 0 reads `flipped`.
    packed = (
        (int(round(float(tensor[0, target_pixel].item()) * 255)) << 16)
        | (int(round(float(tensor[1, target_pixel].item()) * 255)) << 8)
        | int(round(float(tensor[2, target_pixel].item()) * 255))
    )
    mask = 0x3 << (2 * 11)
    packed = (packed & ~mask) | ("ACGT".index(flipped) << (2 * 11))
    tensor[0, target_pixel] = ((packed >> 16) & 0xFF) / 255.0
    tensor[1, target_pixel] = ((packed >> 8) & 0xFF) / 255.0
    tensor[2, target_pixel] = (packed & 0xFF) / 255.0

    _, disagreement = strip_to_sequence(tensor, pixel_stride_bp=6, sequence_length=len(sequence))
    affected_position = offsets[target_pixel] + 0
    assert disagreement[affected_position] > 0