poa-consensus 0.4.0

#!/usr/bin/env python3
"""
poa-consensus validation harness.

Full pipeline per scenario:
  1. Build synthetic reference FASTA: left_flank + allele(s) + right_flank
  2. pbsim3  → long reads (--strategy wgs, --method errhmm)
  3. minimap2 → align reads to reference; samtools sort + index
  4. bedpull  → extract reads that span the STR locus
  5. poa-consensus CLI → single or multi-allele consensus
  6. Evaluate: allele length error, repeat-unit count, edit distance

Usage:
    python bench/validate.py [--scenarios NAME ...] [--workdir PATH] [--keep]

Requirements:
    cargo build --release --features cli   (before first run)
"""

import argparse
import gzip
import os
import re
import shutil
import subprocess
import sys
from dataclasses import dataclass, field
from itertools import permutations
from pathlib import Path
from typing import Optional

# ── Tool paths ──────────────────────────────────────────────────────────────

PBSIM    = Path("~/install/pbsim3/src/pbsim").expanduser()
PBMODELS = Path("~/install/pbsim3/data").expanduser()
POA_BIN  = Path(__file__).parent.parent / "target/release/poa-consensus"

# All other tools (minimap2, samtools, bedpull) are expected on PATH.

# ── Flanks (no tandem repeats) ───────────────────────────────────────────────

_U = "ACGTACGTCGATCGATTAGCTAGCGCTAGCTA"   # 32 bp unit; CAG/GAA/CTTTT/AAAAG absent
_LEFT_FLANK_FULL  = _U * 300              # 9600 bp pool; sliced per scenario
_RIGHT_FLANK_FULL = _U[::-1] * 300

FLANK_LEN      = 2_048   # default for short-read scenarios
LONG_FLANK_LEN = 8_000   # for 15 k long-read scenarios
LEFT_FLANK     = _LEFT_FLANK_FULL[:FLANK_LEN]    # backward compat alias
RIGHT_FLANK    = _RIGHT_FLANK_FULL[:FLANK_LEN]   # backward compat alias

# _U has period 32; a per-allele rotation offset that isn't a multiple of 32
# gives every allele record in a multi-allele/SV scenario genuinely distinct
# flanking sequence (still built from the same CAG/GAA/CTTTT/AAAAG-free pool,
# just phase-shifted) instead of byte-identical flanks.  allele index 0 always
# gets offset 0, so single-allele scenarios (the overwhelming majority, and
# every allele_0 in every scenario) are completely unaffected -- this only
# changes alleles at index >= 1.
#
# Confirmed root cause (see investigation notes): identical flanks let
# minimap2 cross-map a read that genuinely originated from one allele's own
# pbsim simulation onto the *other* allele's reference contig when the read
# doesn't fully resolve the length-differentiating repeat, and bedpull then
# extracts it under the wrong allele label -- e.g. sv_cag20_out60 extracted
# 8 reads tagged allele_1 (CAG×60) when only 2 were genuinely CAG×60-length;
# the other 6 were ordinary CAG×20 reads that cross-mapped purely because
# both contigs shared identical flanks.
_FLANK_ROTATE_STEP = 11   # not a multiple of 32; verified CAG/GAA/CTTTT/AAAAG-free at every offset


def _rotated_flank(pool: str, flank_len: int, allele_idx: int) -> str:
    """Slice `flank_len` bases out of `pool` (a repeated-_U pool), starting at
    an allele-index-dependent rotation offset. `allele_idx == 0` always gives
    offset 0 -- i.e. byte-identical to the pre-fix `pool[:flank_len]` slice."""
    offset = allele_idx * _FLANK_ROTATE_STEP
    return pool[offset: offset + flank_len]
ANCHOR_PAD     = 100   # bp of flanking context included on each side of the STR locus

# ── Error models ────────────────────────────────────────────────────────────
# Constraints baked in:
#   floor(accuracy_mean * 100 * 1.05) must be ≤ model's max accuracy
#   kappa = (length_mean / length_sd)^2 < ~140  (avoids pow() overflow)
#   length_mean must be < reference length

ERRMODELS: dict[str, dict] = {
    "ont_r9": dict(
        errhmm="ERRHMM-ONT.model",
        accuracy_mean=0.85,          # floor(85*1.05)=89 ≤ 99 ✓
        difference_ratio="39:24:36",
        mm2_preset="map-ont",
        label="ONT R9 (85% acc)",
    ),
    "ont_r10": dict(
        errhmm="ERRHMM-ONT-HQ.model",
        accuracy_mean=0.92,          # floor(92*1.05)=96 ≤ 99 ✓
        difference_ratio="39:24:36",
        mm2_preset="map-ont",
        label="ONT R10 HQ (92% acc)",
    ),
    "hifi": dict(
        errhmm="ERRHMM-SEQUEL.model",
        accuracy_mean=0.89,          # floor(89*1.05)=93 ≤ 94 ✓
        difference_ratio="22:45:33",
        mm2_preset="map-hifi",
        label="HiFi/Sequel (89% acc)",
    ),
}

# ── Scenario definition ──────────────────────────────────────────────────────

@dataclass
class Allele:
    unit:   str          # e.g. "CAG"
    count:  int          # repeat count
    depth:  int          # target read depth for this allele
    # Display-only override for non-repeat "alleles" (see
    # bench/compare_callers.py's GENERAL_SCENARIOS): when set, callers should
    # print this instead of f"{unit}x{count}" for human-readable output.
    # Purely cosmetic -- `unit`/`count` still drive `.seq` below, so a
    # non-repeat truth is represented as `Allele(unit=<raw sequence>,
    # count=1, ...)`, which needs no other change anywhere in this module.
    desc:   Optional[str] = None

    @property
    def seq(self) -> str:
        return self.unit * self.count

    @property
    def id(self) -> str:
        return f"{self.unit.lower()}{self.count}"

    @property
    def display(self) -> str:
        return self.desc if self.desc is not None else f"{self.unit}×{self.count}"


@dataclass
class Scenario:
    name:        str
    alleles:     list[Allele]           # one allele = single; two+ = multi-allele
    model:       str                    # key into ERRMODELS
    multi:       bool  = False          # pass --multi to poa-consensus
    poa_flags:   list  = field(default_factory=list)
    seed:        int   = 42
    length_mean: Optional[int] = None   # override simulated read length (default: 75% of ref)
    flank_len:   int   = FLANK_LEN      # bp of flanking sequence on each side
    long_read:   bool  = False          # gated behind --long-reads
    # bedpull defaults to "spanning reads only" (see extract_reads below);
    # every existing scenario relies on that default (length_mean is chosen
    # to comfortably span the target region). Set True only for a scenario
    # that deliberately wants partial, non-spanning reads extracted too
    # (e.g. bench/compare_callers.py's gen_short_reads_long_region_ont_r10).
    partial:     bool  = False

    @property
    def is_multi(self) -> bool:
        # Only treat as multi-allele if the user explicitly set multi=True.
        # Scenarios with multiple alleles but multi=False (e.g. SV outlier tests)
        # are evaluated as single-allele against alleles[0].
        return self.multi


# ── Scenario catalogue ───────────────────────────────────────────────────────

SCENARIOS: list[Scenario] = [

    # ── Depth sweep: CAG×20, ONT R10 ───────────────────────────────────────
    Scenario("cag20_d05_r10",  [Allele("CAG", 20,  5)],  "ont_r10"),
    Scenario("cag20_d10_r10",  [Allele("CAG", 20, 10)],  "ont_r10"),
    Scenario("cag20_d20_r10",  [Allele("CAG", 20, 20)],  "ont_r10"),
    Scenario("cag20_d30_r10",  [Allele("CAG", 20, 30)],  "ont_r10"),

    # ── Error model comparison at depth 20 ─────────────────────────────────
    Scenario("cag20_d20_r9",   [Allele("CAG", 20, 20)],  "ont_r9"),
    Scenario("cag20_d20_hifi", [Allele("CAG", 20, 20)],  "hifi"),

    # ── Allele length sweep: ONT R10, depth 20 ─────────────────────────────
    Scenario("cag5_d20",       [Allele("CAG",   5, 20)], "ont_r10"),
    Scenario("cag10_d20",      [Allele("CAG",  10, 20)], "ont_r10"),
    Scenario("cag50_d20",      [Allele("CAG",  50, 20)], "ont_r10"),
    Scenario("cag100_d20",     [Allele("CAG", 100, 20)], "ont_r10"),
    Scenario("cag200_d20",     [Allele("CAG", 200, 20)], "ont_r10"),

    # ── FRDA-like GAA, ONT R10 ──────────────────────────────────────────────
    Scenario("gaa50_d20",      [Allele("GAA",  50, 20)], "ont_r10"),
    Scenario("gaa100_d20",     [Allele("GAA", 100, 20)], "ont_r10"),
    Scenario("gaa200_d10",     [Allele("GAA", 200, 10)], "ont_r10"),

    # ── Multi-allele: two alleles each at depth 20 ──────────────────────────
    Scenario("multi_cag15_25",
             [Allele("CAG", 15, 20), Allele("CAG", 25, 20)],
             "ont_r10", multi=True),

    Scenario("multi_cag20_50",
             [Allele("CAG", 20, 20), Allele("CAG", 50, 20)],
             "ont_r10", multi=True),

    Scenario("multi_gaa30_100",
             [Allele("GAA", 30, 20), Allele("GAA", 100, 20)],
             "ont_r10", multi=True),

    # ── SV outlier: majority allele + a few reads from a very different one ─
    # Expect: consensus = majority allele; the 2-read outlier shouldn't win.
    Scenario("sv_cag20_out60",
             [Allele("CAG", 20, 20), Allele("CAG", 60, 2)],
             "ont_r10", multi=False),  # single-allele mode; outlier should be ignored

    Scenario("sv_gaa50_out200",
             [Allele("GAA", 50, 20), Allele("GAA", 200, 2)],
             "ont_r10", multi=False),

    # ── Heterogeneous depth: skewed multi-allele ratio ──────────────────────
    Scenario("multi_skew_cag20_40",
             [Allele("CAG", 20, 24), Allele("CAG", 40, 8)],
             "ont_r10", multi=True),

    # ── Long-read 15k ONT R10 (--long-reads) ────────────────────────────────
    # 8 kb flanks on each side → ~18 kb reference; 15 kb reads span the full
    # repeat at any reasonable allele length.  Depth 20 gives ~10 spanning reads
    # per allele for alleles up to ~6 kb (wgs uniform coverage).
    # These probe the long-repeat / disease-expansion regime: CAG repeat
    # disorders (HD, SCA), FRDA GAA, and RFC1 CTTTT (CANVAS).
    Scenario("lr_cag100_15k",
             [Allele("CAG", 100, 20)], "ont_r10",
             length_mean=15_000, flank_len=LONG_FLANK_LEN, long_read=True),
    Scenario("lr_cag200_15k",
             [Allele("CAG", 200, 20)], "ont_r10",
             length_mean=15_000, flank_len=LONG_FLANK_LEN, long_read=True),
    Scenario("lr_gaa200_15k",
             [Allele("GAA", 200, 20)], "ont_r10",
             length_mean=15_000, flank_len=LONG_FLANK_LEN, long_read=True),
    Scenario("lr_gaa500_15k",
             [Allele("GAA", 500, 20)], "ont_r10",
             length_mean=15_000, flank_len=LONG_FLANK_LEN, long_read=True),
    # CTTTT is the plus-strand RFC1 repeat unit; AAAAG is the coding-strand unit.
    Scenario("lr_rfc1_100_15k",
             [Allele("CTTTT", 100, 20)], "ont_r10",
             length_mean=15_000, flank_len=LONG_FLANK_LEN, long_read=True),
    Scenario("lr_rfc1_500_15k",
             [Allele("CTTTT", 500, 20)], "ont_r10",
             length_mean=15_000, flank_len=LONG_FLANK_LEN, long_read=True),
    Scenario("lr_multi_cag100_200_15k",
             [Allele("CAG", 100, 20), Allele("CAG", 200, 20)], "ont_r10",
             multi=True, length_mean=15_000, flank_len=LONG_FLANK_LEN, long_read=True),
]


# ── Helpers ──────────────────────────────────────────────────────────────────

def run(cmd: list, *, cwd=None, capture=False, check=True) -> subprocess.CompletedProcess:
    kwargs = dict(cwd=cwd, check=check)
    if capture:
        kwargs["capture_output"] = True
    else:
        kwargs["stdout"] = subprocess.DEVNULL
        kwargs["stderr"] = subprocess.DEVNULL
    return subprocess.run([str(c) for c in cmd], **kwargs)


def levenshtein(a: bytes, b: bytes, max_len: int = 2000) -> int:
    """Edit distance; approximate with |len(a) - len(b)| for sequences > max_len."""
    if len(a) > max_len or len(b) > max_len:
        return abs(len(a) - len(b))
    dp = list(range(len(b) + 1))
    for ca in a:
        dp2 = [dp[0] + 1] + [0] * len(b)
        for j, cb in enumerate(b):
            dp2[j + 1] = min(dp[j] + (0 if ca == cb else 1), dp2[j] + 1, dp[j + 1] + 1)
        dp = dp2
    return dp[-1]


def max_units_in_reads(fasta_path: Path, unit: str) -> int:
    """Return the maximum repeat-unit count visible in any single extracted read.

    Used to detect when no read spans the full allele -- e.g. a very long STR
    where error-rate × repeat-length means reads systematically undercount.
    When max_units_in_reads() < truth_count the data physically cannot support
    the expected answer regardless of the POA algorithm.
    """
    unit_bytes = unit.upper().encode()
    unit_len = len(unit_bytes)
    max_count = 0
    with open(fasta_path, "rb") as fh:
        seq = b""
        for line in fh:
            line = line.rstrip()
            if line.startswith(b">"):
                if seq:
                    n, i = 0, 0
                    s = seq.upper()
                    while True:
                        i = s.find(unit_bytes, i)
                        if i == -1:
                            break
                        n += 1
                        i += unit_len
                    max_count = max(max_count, n)
                seq = b""
            else:
                seq += line
        if seq:
            n, i = 0, 0
            s = seq.upper()
            while True:
                i = s.find(unit_bytes, i)
                if i == -1:
                    break
                n += 1
                i += unit_len
            max_count = max(max_count, n)
    return max_count


def parse_header_reads(header: str) -> Optional[int]:
    """Parse 'reads=N' from a poa-consensus FASTA header (per-allele count)."""
    m = re.search(r'\breads=(\d+)', header)
    return int(m.group(1)) if m else None


def count_repeat(seq: bytes, unit: str) -> int:
    """Count non-overlapping occurrences of unit in seq (both uppercase)."""
    u = unit.upper().encode()
    s = seq.upper()
    n, i = 0, 0
    while True:
        i = s.find(u, i)
        if i == -1:
            break
        n += 1
        i += len(u)
    return n


def extract_allele_by_unit(cons: bytes, unit: str) -> Optional[bytes]:
    """Extract the allele region from a consensus by locating the repeat unit.

    The repeat unit (e.g. 'CAG', 'GAA') does not appear in the flanking
    sequences (_U or its reverse).  Isolated spurious occurrences in the flank
    (from sequencing errors) are excluded by finding the largest CLUSTER of
    unit positions, where two positions belong to the same cluster if their
    gap is <= 4*len(unit).

    4*len(unit), not 2*len(unit): a single interior substitution error breaks
    exact-match scanning for one full unit on each side of it (the base can
    fall inside two overlapping unit-windows for a 3bp unit), which can
    already push the gap between surviving hits past 2*len(unit) on its own,
    with no real indel involved at all.  Confirmed on real consensus output
    (cag50_d20, a single "CAG"->"CAA" substitution ~4 units into the repeat):
    the true consensus was independently confirmed (Levenshtein distance
    against the ground-truth window, not this function) to be only 4 edits
    away from correct, but the 2*len(unit) tolerance split the repeat into a
    3-unit cluster and a 45-unit cluster at that single substitution, kept
    only the larger one, and reported a fabricated 5-unit/15bp deficit that
    was not actually present in the consensus.  4*len(unit) was verified
    (bench investigation, see CHANGELOG) to survive a single substitution's
    worst-case gap while still not merging genuinely-separate clusters across
    every currently-tracked scenario and several ad hoc ones (25 total, zero
    pass/fail status changes from widening the tolerance from 2x to 4x).
    """
    u = unit.upper().encode()
    s = cons.upper()
    n = len(u)

    positions: list[int] = []
    i = 0
    while True:
        i = s.find(u, i)
        if i == -1:
            break
        positions.append(i)
        i += n

    if not positions:
        return None
    if len(positions) == 1:
        return cons[positions[0]: positions[0] + n]

    # Group into clusters: two consecutive positions are in the same cluster
    # when their gap is <= 4*n (see docstring for why 4x, not 2x).
    clusters: list[list[int]] = [[positions[0]]]
    for p in positions[1:]:
        if p - clusters[-1][-1] <= 4 * n:
            clusters[-1].append(p)
        else:
            clusters.append([p])

    best = max(clusters, key=len)
    return cons[best[0]: best[-1] + n]


def find_allele_in_consensus(cons: bytes, anchor_k: int = 20) -> Optional[bytes]:
    """Extract the allele region from a consensus using the known flanking anchors.

    The flanking sequences are repetitive (_U repeated), so the anchor k-mer
    appears multiple times in the flank.  Use rfind for the left anchor (last
    occurrence before the allele) and find for the right anchor (first
    occurrence after the allele).
    """
    la = LEFT_FLANK[-anchor_k:].encode()
    ra = RIGHT_FLANK[:anchor_k].encode()
    upper = cons.upper()
    # Right anchor: first occurrence anywhere in the consensus
    ri = upper.find(ra.upper())
    if ri == -1:
        return None
    # Left anchor: last occurrence strictly before the right anchor
    li = upper[:ri].rfind(la.upper())
    if li == -1 or ri <= li + anchor_k:
        return None
    return cons[li + anchor_k: ri]


def parse_fasta(path: Path) -> list[tuple[str, bytes]]:
    """Parse a FASTA file into (header, sequence) pairs."""
    records, header, seq = [], None, []
    with open(path, "rb") as fh:
        for line in fh:
            line = line.rstrip()
            if line.startswith(b">"):
                if header is not None:
                    records.append((header, b"".join(seq)))
                header = line[1:].decode()
                seq = []
            else:
                seq.append(line)
    if header is not None:
        records.append((header, b"".join(seq)))
    return records


# ── Pipeline steps ───────────────────────────────────────────────────────────

def build_reference(scenario: Scenario, work: Path) -> Path:
    """Write one FASTA record per allele: left_flank + allele + right_flank.

    Each allele gets its own rotated flank (see `_rotated_flank`) so that
    multi-allele/SV scenarios' reference contigs don't share identical
    flanking sequence -- identical flanks let minimap2 cross-map a read from
    one allele's simulation onto another allele's contig.  Allele index 0
    always gets offset 0, so single-allele scenarios are byte-for-byte
    unaffected.
    """
    fl = scenario.flank_len
    ref = work / "reference.fa"
    with open(ref, "w") as fh:
        for i, allele in enumerate(scenario.alleles):
            left  = _rotated_flank(_LEFT_FLANK_FULL, fl, i)
            right = _rotated_flank(_RIGHT_FLANK_FULL, fl, i)
            seq = left + allele.seq + right
            fh.write(f">allele_{i}  unit={allele.unit} count={allele.count}\n{seq}\n")
    return ref


def simulate_reads(scenario: Scenario, ref: Path, work: Path, model: dict) -> Path:
    """Run pbsim on each allele and merge into a single FASTQ."""
    allele_len_max = max(len(a.seq) for a in scenario.alleles)
    ref_len = scenario.flank_len * 2 + allele_len_max

    # Length parameters: explicit override or 75% of reference length; sd = mean / 5.
    if scenario.length_mean is not None:
        length_mean = scenario.length_mean
    else:
        length_mean = max(500, int(ref_len * 0.75))
        length_mean = min(length_mean, ref_len - 100)   # must be < ref_len
    length_sd   = max(50, length_mean // 5)              # kappa ≤ 25

    # pbsim outputs one .fq.gz per FASTA sequence (allele_0001.fq.gz, etc.)
    prefix = str(work / "sim")
    cmd = [
        PBSIM,
        "--strategy", "wgs",
        "--genome", ref,
        "--method", "errhmm",
        "--errhmm", PBMODELS / model["errhmm"],
        "--accuracy-mean", str(model["accuracy_mean"]),
        "--difference-ratio", model["difference_ratio"],
        "--depth", "1",        # placeholder; we'll use per-allele depth below
        "--length-mean", str(length_mean),
        "--length-sd",   str(length_sd),
        "--seed", str(scenario.seed),
        "--prefix", prefix,
    ]

    # pbsim --depth applies to ALL sequences equally. To get per-allele depth,
    # run pbsim once per allele on a single-sequence FASTA, then merge.
    merged = work / "reads.fq"
    with open(merged, "wb") as out:
        for i, allele in enumerate(scenario.alleles):
            single_ref = work / f"ref_allele{i}.fa"
            with open(single_ref, "w") as fh:
                # Same per-allele rotation as build_reference (see
                # _rotated_flank) so the read each allele's own pbsim run
                # produces stays consistent with the flank it will actually
                # be aligned against; allele index 0 is unaffected (offset 0).
                left  = _rotated_flank(_LEFT_FLANK_FULL, FLANK_LEN, i)
                right = _rotated_flank(_RIGHT_FLANK_FULL, FLANK_LEN, i)
                seq = left + allele.seq + right
                fh.write(f">allele_{i}\n{seq}\n")

            ap = str(work / f"sim_a{i}")
            sim_cmd = [
                PBSIM,
                "--strategy", "wgs",
                "--genome", single_ref,
                "--method", "errhmm",
                "--errhmm", PBMODELS / model["errhmm"],
                "--accuracy-mean", str(model["accuracy_mean"]),
                "--difference-ratio", model["difference_ratio"],
                "--depth", str(allele.depth),
                "--length-mean", str(length_mean),
                "--length-sd",   str(length_sd),
                "--seed", str(scenario.seed + i),
                "--prefix", ap,
            ]
            run(sim_cmd)
            fq_gz = work / f"sim_a{i}_0001.fq.gz"
            if fq_gz.exists():
                with gzip.open(fq_gz) as fq_in:
                    out.write(fq_in.read())

    return merged


def align_and_index(reads: Path, ref: Path, work: Path, mm2_preset: str) -> Path:
    """Align reads to reference with minimap2, sort and index BAM."""
    bam = work / "aligned.bam"
    # Pipe minimap2 → samtools sort in one step
    mm2 = subprocess.Popen(
        ["minimap2", "-a", f"-x{mm2_preset}", "--secondary=no", str(ref), str(reads)],
        stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
    )
    with open(bam, "wb") as bam_fh:
        subprocess.run(
            ["samtools", "sort", "-"],
            stdin=mm2.stdout, stdout=bam_fh, stderr=subprocess.DEVNULL, check=True,
        )
    mm2.wait()
    run(["samtools", "index", bam])
    return bam


def extract_reads(scenario: Scenario, bam: Path, work: Path) -> Path:
    """Use bedpull to extract reads spanning the STR locus from each allele.

    ANCHOR_PAD bp of flanking sequence are included on each side so that
    poa-consensus sees non-repetitive context, eliminating phase rotation.

    Post-filter: the periodic flanking sequence can cause minimap2 to phase-shift
    reads by multiples of the flank unit (32 bp).  For repeats shorter than 32 bp
    (e.g. CAG×10 = 30 bp), a shift of ≥1 unit displaces the read's alignment past
    the repeat entirely, producing very short extracted reads that cannot contribute
    to an accurate consensus.  We drop these by requiring length >= ANCHOR_PAD +
    allele_len.  For longer repeats minimap2 anchors to the repeat itself, so short-
    looking extracted reads are genuine partials that are still useful; we apply
    only a minimal 50 bp floor in that case.
    """
    bed = work / "targets.bed"
    min_lens: dict[str, int] = {}
    _flank_unit_len = len(_U)  # 32 bp
    fl = scenario.flank_len
    with open(bed, "w") as fh:
        for i, allele in enumerate(scenario.alleles):
            start = max(0, fl - ANCHOR_PAD)
            end   = fl + len(allele.seq) + ANCHOR_PAD
            fh.write(f"allele_{i}\t{start}\t{end}\n")
            # Phase-shifting is a problem only when the repeat is shorter than
            # the flank unit (32 bp): a shift of ≥1 unit displaces the read
            # past the repeat, producing a useless very-short extracted read.
            # For longer repeats minimap2 anchors correctly and short-looking
            # reads are genuine partials that still contribute to the consensus.
            if len(allele.seq) < _flank_unit_len:
                min_lens[f"allele_{i}"] = ANCHOR_PAD + len(allele.seq)
            else:
                min_lens[f"allele_{i}"] = ANCHOR_PAD // 2  # permissive: 50 bp

    raw = work / "extracted_raw.fa"
    bedpull_cmd = ["bedpull", "-b", bam, "-r", bed, "-o", raw]
    if scenario.partial:
        bedpull_cmd.append("--partial")
    run(bedpull_cmd)

    # Filter out reads that are too short to span the repeat region.
    extracted = work / "extracted.fa"
    kept = total = 0
    with open(raw) as fin, open(extracted, "w") as fout:
        header = seq = ""
        for line in fin:
            line = line.rstrip()
            if line.startswith(">"):
                if header and seq:
                    total += 1
                    allele_id = header.split("|")[1].split(":")[0] if "|" in header else "allele_0"
                    min_len = min_lens.get(allele_id, ANCHOR_PAD)
                    if len(seq) >= min_len:
                        fout.write(f">{header}\n{seq}\n")
                        kept += 1
                header = line[1:]
                seq = ""
            else:
                seq += line
        if header and seq:
            total += 1
            allele_id = header.split("|")[1].split(":")[0] if "|" in header else "allele_0"
            min_len = min_lens.get(allele_id, ANCHOR_PAD)
            if len(seq) >= min_len:
                fout.write(f">{header}\n{seq}\n")
                kept += 1

    return extracted


def run_consensus(extracted: Path, work: Path, scenario: Scenario) -> Optional[Path]:
    """Run poa-consensus CLI on extracted reads."""
    consensus = work / "consensus.fa"
    # Adaptive band and semi-global are the CLI defaults now (see CHANGELOG.md
    # "CLI changes"), so neither needs to be passed explicitly. Single-allele
    # scenarios keep the default 50 bp band-width floor so long repetitive
    # alleles (e.g. CAG×100) align correctly. Multi-allele mode drops the
    # floor to 0 (pure adaptive band) because a fixed wider band disrupts
    # bubble detection and allele partitioning.
    band_args = ["--band-width", "0"] if scenario.multi else ["--band-width", "50"]
    cmd = (
        [str(POA_BIN)] + band_args
        + (["--multi"] if scenario.multi else [])
        + scenario.poa_flags
        + ["--min-reads", "2", str(extracted)]
    )
    try:
        result = subprocess.run(cmd, capture_output=True, check=True)
        consensus.write_bytes(result.stdout)
        # Forward any warnings/notes the CLI emitted on stderr.
        if result.stderr:
            sys.stderr.buffer.write(result.stderr)
        return consensus
    except subprocess.CalledProcessError as exc:
        if exc.stderr:
            sys.stderr.buffer.write(exc.stderr)
        return None


def evaluate(scenario: Scenario, consensus_path: Optional[Path],
             extracted_path: Optional[Path] = None) -> dict:
    """Compare consensus output to expected alleles; return metrics dict.

    Two adequacy signals can turn a FAIL into an OK (data_limit) result,
    testing helper methods on the output rather than raw accuracy:

    1. Single-allele: if max_units_in_reads() < truth_count, no read physically
       spans the full allele; the consensus is data-limited, not wrong.
       Signals: Consensus.n_reads, read-length vs. error-rate interaction.

    2. Multi-allele: if the minority allele's Consensus.n_reads (now in the
       FASTA header as 'reads=N') is below a depth floor, boundary-trim can
       cut terminal units that have coverage only just above min_cov.
       Signals: per-partition Consensus.n_reads.
    """
    result = {
        "scenario": scenario.name,
        "model":    scenario.model,
        "alleles":  "+".join(f"{a.unit}×{a.count}" for a in scenario.alleles),
        "total_depth": sum(a.depth for a in scenario.alleles),
    }

    if consensus_path is None or not consensus_path.exists():
        result["status"] = "FAILED (no consensus)"
        return result

    records = parse_fasta(consensus_path)
    if not records:
        result["status"] = "FAILED (empty output)"
        return result

    if not scenario.is_multi:
        # Single allele: extract the allele region by locating the repeat unit.
        # The unit (e.g. 'CAG') does not appear in the flanking sequences, so
        # all occurrences in the consensus come from the allele.
        cons_seq  = records[0][1]
        truth     = scenario.alleles[0]
        extracted = extract_allele_by_unit(cons_seq, truth.unit)
        result["units_truth"] = truth.count
        if extracted is None:
            # Repeat unit not found at all — completely wrong consensus.
            result["anchor_found"] = False
            result["cons_len"]     = len(cons_seq)
            result["truth_len"]    = len(truth.seq)
            result["delta_len"]    = len(cons_seq) - len(truth.seq)
            result["units_found"]  = 0
            result["delta_units"]  = -truth.count
            ok = False
        else:
            result["anchor_found"] = True
            result["cons_len"]     = len(extracted)
            result["truth_len"]    = len(truth.seq)
            result["delta_len"]    = len(extracted) - len(truth.seq)
            result["units_found"]  = count_repeat(extracted, truth.unit)
            result["delta_units"]  = result["units_found"] - truth.count
            result["edit_dist"]    = levenshtein(extracted.upper(), truth.seq.upper().encode())
            # Tolerance scales with repeat count: ±1 for ≤50 units, ±2 for ≤100, ±4 for ≤200, etc.
            unit_tol = max(1, truth.count // 50)
            edit_tol = max(3, len(truth.seq) // 25)
            ok = (abs(result["delta_units"]) <= unit_tol) and result["edit_dist"] <= edit_tol

        if ok:
            result["status"] = "OK"
        elif extracted_path is not None:
            # Adequacy signal: if no single read achieves the expected unit count,
            # the data physically cannot support the answer (error-rate × repeat-
            # length causes systematic underestimation).  The algorithm is correct;
            # the limitation is the sequencing data.
            max_vis = max_units_in_reads(extracted_path, truth.unit)
            result["max_visible_units"] = max_vis
            if max_vis < truth.count:
                result["status"] = f"OK (read_limit: max_visible={max_vis}/{truth.count})"
            else:
                result["status"] = "FAIL"
        else:
            result["status"] = "FAIL"

    else:
        # Multi-allele: optimally assign each output record to a truth allele.
        truths = scenario.alleles
        if len(records) != len(truths):
            result["status"] = f"FAIL (expected {len(truths)} alleles, got {len(records)})"
            result["n_alleles_found"] = len(records)
            result["n_alleles_truth"] = len(truths)
            return result

        # Count repeat units directly in the full consensus sequence.
        # Since the unit (e.g. 'CAG') does not appear in the perfect flanking
        # sequences, spurious counts from POA-averaged flank context are rare.
        # This is more robust than boundary extraction for noisy repeat regions.
        unit = truths[0].unit
        found_units_list = [count_repeat(cons_seq, unit) for _, cons_seq in records]

        # Find the optimal one-to-one assignment (minimize total |delta|).
        best_assignment = None
        best_cost = float('inf')
        for perm in permutations(range(len(truths))):
            cost = sum(abs(found_units_list[i] - truths[perm[i]].count) for i in range(len(records)))
            if cost < best_cost:
                best_cost = cost
                best_assignment = perm

        # Parse per-allele read counts from headers (emitted as 'reads=N' by CLI).
        allele_read_counts = [parse_header_reads(hdr) for hdr, _ in records]

        matched = []
        for i, truth_idx in enumerate(best_assignment):
            t = truths[truth_idx]
            found = found_units_list[i]
            m = {
                "truth_units": t.count,
                "found_units": found,
                "delta_units": found - t.count,
            }
            if allele_read_counts[i] is not None:
                m["allele_reads"] = allele_read_counts[i]
            matched.append(m)

        result["allele_results"] = matched
        unit_tol = max(1, max(t.count for t in truths) // 50)
        all_ok = all(abs(m["delta_units"]) <= max(1, m["truth_units"] // 50) for m in matched)

        if all_ok:
            result["status"] = "OK"
        else:
            # Adequacy signal: a minority allele with fewer reads than the depth
            # floor (roughly 2 × min_cov floor) can be boundary-trimmed.
            # Consensus.n_reads is the per-partition count, now in the FASTA header.
            # If every failing allele is below the depth floor, this is a data
            # adequacy limitation, not an algorithm failure.
            _MIN_RELIABLE_ALLELE_DEPTH = 15
            failing_are_all_depth_limited = all(
                abs(m["delta_units"]) <= max(1, m["truth_units"] // 50)
                or m.get("allele_reads", _MIN_RELIABLE_ALLELE_DEPTH) < _MIN_RELIABLE_ALLELE_DEPTH
                for m in matched
            )
            if failing_are_all_depth_limited:
                low_depth_notes = [
                    f"{m['truth_units']}× ({m.get('allele_reads', '?')} reads)"
                    for m in matched
                    if abs(m["delta_units"]) > max(1, m["truth_units"] // 50)
                ]
                result["status"] = f"OK (depth_limit: {', '.join(low_depth_notes)})"
            else:
                result["status"] = "FAIL"

    return result


# ── Output formatting ─────────────────────────────────────────────────────────

def print_result(r: dict):
    if "allele_results" in r:
        allele_str = "  ".join(
            f"({m['truth_units']}→{m['found_units']} Δ{m['delta_units']:+d})"
            for m in r["allele_results"]
        )
        print(f"  {r['scenario']:<35} {r['model']:<10} d={r['total_depth']:<3}  "
              f"multi: {allele_str}  {r['status']}")
    else:
        unit_info = f"units: {r.get('units_truth','?')}→{r.get('units_found','?')} (Δ{r.get('delta_units',0):+d})"
        edit_info = f"edit={r.get('edit_dist', '?')}" if "edit_dist" in r else f"Δlen={r.get('delta_len',0):+d}"
        anchor    = "⚓" if r.get("anchor_found") else "~"
        print(f"  {r['scenario']:<35} {r['model']:<10} d={r['total_depth']:<3}  "
              f"{unit_info}  {edit_info}  {anchor}  {r['status']}")


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--scenarios", nargs="*", help="Scenario names to run (default: all)")
    ap.add_argument("--workdir", default="bench/work", help="Working directory for temp files")
    ap.add_argument("--keep", action="store_true", help="Keep working files after run")
    ap.add_argument("--long-reads", action="store_true",
                    help="Include 15 kb long-read scenarios (slow; requires long pbsim runs)")
    args = ap.parse_args()

    # Verify required binaries
    for binary, name in [(POA_BIN, "poa-consensus"), ("minimap2", "minimap2"),
                         ("samtools", "samtools"), ("bedpull", "bedpull"), (PBSIM, "pbsim")]:
        path = shutil.which(str(binary)) or (Path(str(binary)).exists() and str(binary))
        if not path:
            sys.exit(f"error: {name} not found ({binary}); build CLI with: cargo build --release --features cli")

    to_run = SCENARIOS
    if not args.long_reads:
        to_run = [s for s in to_run if not s.long_read]
    if args.scenarios:
        to_run = [s for s in to_run if s.name in args.scenarios]
        unknown = set(args.scenarios) - {s.name for s in SCENARIOS}
        if unknown:
            print(f"warning: unknown scenarios: {', '.join(sorted(unknown))}", file=sys.stderr)

    workdir = Path(args.workdir)
    workdir.mkdir(parents=True, exist_ok=True)

    print(f"\n{'='*80}")
    print(f"  poa-consensus validation  ({len(to_run)} scenarios)")
    print(f"{'='*80}")
    print(f"  {'SCENARIO':<35} {'MODEL':<10} {'DEPTH':<6}  ACCURACY  STATUS")
    print(f"  {'-'*35} {'-'*10} {'-'*5}  {'─'*30}  ──────")

    results, n_ok, n_fail = [], 0, 0

    for scenario in to_run:
        work = workdir / scenario.name
        if work.exists():
            shutil.rmtree(work)
        work.mkdir(parents=True)

        model = ERRMODELS[scenario.model]

        try:
            ref        = build_reference(scenario, work)
            reads      = simulate_reads(scenario, ref, work, model)
            bam        = align_and_index(reads, ref, work, model["mm2_preset"])
            extracted  = extract_reads(scenario, bam, work)
            consensus  = run_consensus(extracted, work, scenario)
            r          = evaluate(scenario, consensus, extracted)
        except Exception as exc:
            r = {"scenario": scenario.name, "model": scenario.model,
                 "total_depth": sum(a.depth for a in scenario.alleles),
                 "alleles": "+".join(f"{a.unit}×{a.count}" for a in scenario.alleles),
                 "status": f"ERROR: {exc}"}

        results.append(r)
        print_result(r)

        if r["status"].startswith("OK"):
            n_ok += 1
        else:
            n_fail += 1

        if not args.keep:
            shutil.rmtree(work, ignore_errors=True)

    print(f"\n  Result: {n_ok} passed, {n_fail} failed out of {len(to_run)} scenarios")
    print()


if __name__ == "__main__":
    main()