tests/test_cli.py

from __future__ import annotations

import json

import pandas as pd

from deepgenopix.cli import (
    build_parser,
    build_experiment_matrix,
    main,
    list_presets,
    parse_overrides,
    runnable_preset_names,
)
from deepgenopix.notebook_support import discover_raw_inputs, discover_split_parquet_root


def test_parse_overrides_requires_json_object():
    try:
        parse_overrides('["baseline_v1"]')
    except ValueError as exc:
        assert "JSON object" in str(exc)
    else:
        raise AssertionError("expected ValueError")


def test_list_presets_excludes_blocked_by_default():
    preset_names = {row["preset"] for row in list_presets()}

    assert "baseline_v1" in preset_names
    assert "zero_flank_v1" not in preset_names
    assert "walk9_v1" in preset_names
    assert "walk6_v1" in preset_names


def test_runnable_preset_names_include_walk_presets():
    assert runnable_preset_names() == [
        "baseline_v1",
        "stride2_v1",
        "stride8_v1",
        "latent128_v1",
        "latent768_v1",
        "layers4_v1",
        "stem5_v1",
        "stem7_v1",
        "walk9_v1",
        "walk6_v1",
    ]


def test_build_experiment_matrix_prefers_recovered_parquet(tmp_path):
    raw_dir = tmp_path / "data" / "raw"
    raw_dir.mkdir(parents=True)
    (raw_dir / "te_seqdata.parquet").write_bytes(b"")
    (raw_dir / "te_seqdata.recovered.parquet").write_bytes(b"")

    matrix = build_experiment_matrix(repo_root=tmp_path)

    assert [row["preset"] for row in matrix] == runnable_preset_names()
    assert matrix[0]["compare_against"] is None
    assert matrix[1]["compare_against"] == "baseline_v1"
    assert matrix[2]["bp_per_token"] == 96
    assert matrix[0]["raw_parquet"] == "data/raw/te_seqdata.recovered.parquet"
    assert "recovered.parquet" in matrix[0]["raw_input_policy"]


def test_build_experiment_matrix_prefers_split_dataset(tmp_path):
    raw_dir = tmp_path / "data" / "raw"
    for split in ("train", "val", "test"):
        split_dir = raw_dir / split
        split_dir.mkdir(parents=True, exist_ok=True)
        pd.DataFrame({"sequence": ["A" * 12], "family": ["fam_a"]}).to_parquet(
            split_dir / "te_seqdata.parquet",
            index=False,
        )
    (raw_dir / "split_summary.json").write_text(json.dumps({"counts": {"train": 1, "val": 1, "test": 1}}), encoding="utf-8")

    matrix = build_experiment_matrix(repo_root=tmp_path)

    assert matrix[0]["raw_parquet"] is None
    assert matrix[0]["raw_split_root"] == "data/raw"
    assert matrix[0]["raw_split_summary"] == "data/raw/split_summary.json"
    assert "split_summary.json" in matrix[0]["raw_input_policy"]


def test_discover_raw_inputs_prefers_recovered_parquet(tmp_path):
    raw_dir = tmp_path / "data" / "raw"
    raw_dir.mkdir(parents=True)
    (raw_dir / "te_seqdata_with_biostats.parquet").write_bytes(b"")
    (raw_dir / "te_seqdata.parquet").write_bytes(b"")
    (raw_dir / "te_seqdata.recovered.parquet").write_bytes(b"")

    raw_fasta, raw_parquet = discover_raw_inputs(raw_dir)

    assert raw_fasta is None
    assert raw_parquet == raw_dir / "te_seqdata.recovered.parquet"


def test_discover_split_parquet_root_returns_summary_path(tmp_path):
    raw_dir = tmp_path / "data" / "raw"
    for split in ("train", "val", "test"):
        split_dir = raw_dir / split
        split_dir.mkdir(parents=True, exist_ok=True)
        pd.DataFrame({"sequence": ["A" * 12], "family": ["fam_a"]}).to_parquet(
            split_dir / "te_seqdata.parquet",
            index=False,
        )
    summary_path = raw_dir / "split_summary.json"
    summary_path.write_text("{}", encoding="utf-8")

    split_root, split_summary = discover_split_parquet_root(raw_dir)

    assert split_root == raw_dir
    assert split_summary == summary_path


def test_cli_prep_subcommand_writes_split(tmp_path):
    source = tmp_path / "te_seqdata.recovered.parquet"
    output_root = tmp_path / "te_split"
    pd.DataFrame(
        {
            "sequence": ["A" * 12, "C" * 12, "G" * 12, "T" * 12],
            "family": ["fam_a", "fam_a", "fam_b", "fam_b"],
        }
    ).to_parquet(source, index=False)

    exit_code = main(
        [
            "prep",
            "--input",
            str(source),
            "--output-root",
            str(output_root),
            "--json",
        ]
    )

    assert exit_code == 0
    assert (output_root / "split_summary.json").exists()
    assert (output_root / "train" / "te_seqdata.parquet").exists()


def test_cli_etl_subcommand_writes_lmdb(tmp_path):
    split_root = tmp_path / "te_split"
    output_dir = tmp_path / "processed"

    for split, rows in {
        "train": [{"sequence": "A" * 12, "family": "fam_a"}],
        "val": [{"sequence": "C" * 12, "family": "fam_b"}],
        "test": [{"sequence": "G" * 12, "family": "fam_c"}],
    }.items():
        split_dir = split_root / split
        split_dir.mkdir(parents=True, exist_ok=True)
        pd.DataFrame(rows).to_parquet(split_dir / "te_seqdata.parquet", index=False)

    exit_code = main(
        [
            "etl",
            "--split-root",
            str(split_root),
            "--output-dir",
            str(output_dir),
            "--pixel-stride-bp",
            "9",
            "--json",
        ]
    )

    assert exit_code == 0
    assert (output_dir / "tensors.lmdb").exists()
    registry = pd.read_csv(output_dir / "registry.csv")
    assert registry["pixel_stride_bp"].tolist() == [9, 9, 9]


def test_cli_help_includes_quant_and_variant_commands():
    help_text = build_parser().format_help()

    assert "quant-train" in help_text
    assert "quant-validate" in help_text
    assert "quantify" in help_text
    assert "variant" in help_text