Blame for webapp/classifier.py

"""The classification pipeline the web app runs: ``taxonova classify`` itself.

Each request's sequences are written to FASTA and classified by the package's
own ``classification_mode`` with a model this process already holds (see
``package_adapter``), so the TSV is the one ``taxonova classify`` writes for
the same model, input and thresholds. The rows of that TSV are what the user
downloads.

Thresholds: with no confidence and no per-rank threshold the model's
calibrated operating point is used, as on the command line; any threshold
switches to the scaled profile (``taxonova.core.thresholds``). The species
gate and the species-group rule of a model apply in both cases, as on the
command line.

Sequences that fail the SSU rRNA screen (``rrna_prefilter``) never reach the
model and are reported as Unclassified at every rank.
"""
from __future__ import annotations

import logging
import os
import tempfile
from typing import Callable, Dict, List, Optional, Sequence, Tuple

from taxonova.core.thresholds import GRADED_SPECIES

from webapp import package_adapter
from webapp.package_adapter import LoadedModel, RunReport, threshold_argv
from webapp.rrna_prefilter import prefilter_sequences

logger = logging.getLogger(__name__)

#: Canonical rank order and TSV columns of the CLI.
LEVEL_ORDER = list(package_adapter.RANKS)

#: The species confidence of the legacy /classify endpoint when none is sent
#: (its 26.10 default, kept so that old scripts get the same thresholds).
#: The page and /api/jobs send no confidence by default: the model's
#: operating point.
LEGACY_CONFIDENCE = GRADED_SPECIES

#: Sequences classified per model call. Labels do not depend on it (every
#: sequence is classified independently); it sets how often progress is
#: reported. test_parity.py checks that. A multiple of the classifier batch
#: size, so batches line up with the command line.
CHUNK_SIZE = int(os.environ.get("TAXONOVA_WEB_CHUNK", "1024"))

Triple = Tuple[str, str, dict]


def load_model(model_path: str) -> LoadedModel:
    """Load one checkpoint with the package loader."""
    return package_adapter.load_model(model_path)


def cli_command(model_id: str, confidence: Optional[float] = None,
                overrides: Optional[Dict[str, float]] = None) -> str:
    """The command that reproduces a web run on the command line."""
    argv = threshold_argv(confidence, overrides)
    shown = [a if a.startswith("--") else f"{float(a):g}" for a in argv]
    return " ".join(["taxonova classify reads.fasta --model", model_id] + shown)


def run_report(cached: LoadedModel, sequences: Sequence[Triple], confidence: Optional[float] = None,
               overrides: Optional[Dict[str, float]] = None, strand_check: bool = True) -> RunReport:
    """Classify ``sequences`` with a loaded model: the CLI TSV and what was withheld.

    ``strand_check`` False skips the package's orientation check (see
    ``package_adapter.classify``)."""
    tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".fasta", delete=False, encoding="utf-8")
    try:
        for seq_id, seq, _ in sequences:
            tmp.write(f">{seq_id}\n{seq}\n")
        tmp.close()
        return package_adapter.classify(cached, tmp.name, confidence, overrides, strand_check)
    finally:
        try:
            os.unlink(tmp.name)
        except OSError:
            pass


def run_classification(cached: LoadedModel, sequences: Sequence[Triple], confidence: Optional[float] = None,
                       overrides: Optional[Dict[str, float]] = None) -> str:
    """Classify ``sequences`` with a loaded model; return the CLI TSV text."""
    return run_report(cached, sequences, confidence, overrides).tsv


def tsv_header() -> str:
    cols = ["sequence_id"]
    for level in LEVEL_ORDER:
        cols += [level, f"{level}_confidence"]
    return "\t".join(cols)


def unclassified_tsv_row(seq_id: str) -> str:
    return "\t".join([seq_id] + ["Unclassified", "0.0000"] * len(LEVEL_ORDER))


def uniquify_sequence_ids(sequences):
    """Give every sequence a unique internal id, keeping the display names.

    FASTA headers need not be unique and the merge keys on the id, so
    duplicates would otherwise all receive one call.
    Returns (sequences_with_unique_ids, display_names_by_position).
    """
    used = set()
    unique, display = [], []
    for seq_id, seq, tax in sequences:
        display.append(seq_id)
        uid, n = seq_id, 1
        while uid in used:
            n += 1
            uid = f"{seq_id}__dup{n}"
        used.add(uid)
        unique.append((uid, seq, tax))
    return unique, display


def restore_display_ids(tsv_string, results, display):
    """Put the user's sequence names back, by position."""
    if len(results) != len(display):
        raise RuntimeError(f"result/display length mismatch: {len(results)} vs {len(display)}")
    for entry, name in zip(results, display):
        entry["sequence_id"] = name
    lines = tsv_string.splitlines()
    if not lines:
        return tsv_string, results
    body = lines[1:]
    if len(body) != len(display):
        raise RuntimeError(f"TSV/display length mismatch: {len(body)} vs {len(display)}")
    out = [lines[0]]
    for ln, name in zip(body, display):
        cols = ln.split("\t")
        cols[0] = name
        out.append("\t".join(cols))
    return "\n".join(out) + "\n", results


def parse_tsv(tsv_string: str) -> List[dict]:
    """CLI TSV text -> one dict per row (confidences as floats)."""
    lines = [ln for ln in tsv_string.splitlines() if ln != ""]
    if not lines:
        return []
    header = lines[0].split("\t")
    col = {name: i for i, name in enumerate(header)}
    rows = []
    for line in lines[1:]:
        f = line.split("\t")
        entry = {"sequence_id": f[col["sequence_id"]]}
        for level in LEVEL_ORDER:
            entry[level] = f[col[level]] if level in col else "Unclassified"
            try:
                entry[f"{level}_confidence"] = float(f[col[f"{level}_confidence"]])
            except (KeyError, ValueError, IndexError):
                entry[f"{level}_confidence"] = 0.0
        rows.append(entry)
    return rows


ProgressFn = Callable[[str, float], None]


def classify_with_prefilter(cached: LoadedModel, sequences: Sequence[Triple], confidence: Optional[float] = None,
                            overrides: Optional[Dict[str, float]] = None,
                            progress: Optional[ProgressFn] = None,
                            chunk_size: int = CHUNK_SIZE) -> Tuple[str, List[dict], dict]:
    """SSU screen, then classify what passed, then merge in input order.

    ``sequences`` must already carry unique ids. Returns the TSV text (CLI
    schema; passing rows byte-identical to the CLI), one dict per input
    sequence with a ``prefilter_pass`` flag, and the run summary
    (``run_info``).
    """
    report = progress or (lambda stage, frac: None)
    report("Screening for 16S rRNA", 0.02)
    pass_ids = prefilter_sequences(list(sequences))
    passing = [t for t in sequences if t[0] in pass_ids]
    logger.info("rRNA screen: %d of %d passed", len(passing), len(sequences))

    header = tsv_header()
    row_by_id: Dict[str, str] = {}
    reports: List[RunReport] = []
    done = 0
    for start, stop in chunk_bounds(len(passing), chunk_size):
        chunk = passing[start:stop]
        report(f"Classifying sequences ({done:,} of {len(passing):,})",
               0.15 + 0.8 * done / max(1, len(passing)))
        # orientation: checked once per job, on the first reads, as the CLI checks a file
        rep = run_report(cached, chunk, confidence, overrides, strand_check=start == 0)
        reports.append(rep)
        lines = [ln for ln in rep.tsv.splitlines() if ln]
        if lines:
            header = lines[0]
            sid_idx = header.split("\t").index("sequence_id")
            for ln in lines[1:]:
                row_by_id[ln.split("\t")[sid_idx]] = ln
        done += len(chunk)
    report("Preparing results", 0.97)

    tsv_lines = [header]
    for seq_id, _, _ in sequences:
        tsv_lines.append(row_by_id.get(seq_id) or unclassified_tsv_row(seq_id))
    tsv = "\n".join(tsv_lines) + "\n"
    rows = parse_tsv(tsv)
    for row, (seq_id, _, _) in zip(rows, sequences):
        row["prefilter_pass"] = seq_id in row_by_id
    return tsv, rows, run_info(cached, reports, passing, confidence, overrides)


def chunk_bounds(n: int, chunk_size: int) -> List[Tuple[int, int]]:
    """(start, stop) of each chunk. The first chunk is made of whole chunks
    enough to hold the reads the package checks for orientation, so that the
    one check per job sees the same reads as the command line."""
    size = max(1, chunk_size)
    sample = package_adapter.strand_sample_size()
    first = min(n, -(-sample // size) * size)
    bounds = [(0, first)] if first else []
    bounds += [(s, min(s + size, n)) for s in range(first, n, size)]
    return bounds


def run_info(cached: LoadedModel, reports: Sequence[RunReport], passing: Sequence[Triple],
             confidence: Optional[float], overrides: Optional[Dict[str, float]]) -> dict:
    """What the page says about a run: thresholds and their source, and why species are missing."""
    settings = getattr(cached, "settings", None) or {}
    first = reports[0] if reports else None
    if first is not None:
        thresholds, source = first.thresholds, first.source
        gate, groups = first.species_gate, first.species_groups
    else:  # nothing passed the screen: what would have applied
        thresholds = None
        source = "user" if (confidence is not None or overrides) else (
            "checkpoint" if settings.get("operating_point") else "default")
        gate, groups = bool(settings.get("species_gate")), settings.get("species_groups", "kept")
    return {
        "thresholds": thresholds,
        "source": source,
        "species_gate": gate,
        "species_groups": groups,
        "gated": sum(len(r.gated) for r in reports),
        "collapsed": sum(len(r.collapsed) for r in reports),
        "low_information": sum(len(r.low_information) for r in reports),
        # classified, with fewer than SHORT_READ_BASES valid bases
        "short_read": sum(len(r.short_read) for r in reports),
        # species confidence at or above the threshold, species Unclassified for another reason
        "higher_rank": sum(len(r.higher_rank) for r in reports),
        "rounding": sum(len(r.rounding) for r in reports),
        "other_rule": sum(len(r.other_rule) for r in reports),
        # per rank: reads at or above the threshold without a label, by reason
        "withheld_by_rank": withheld_by_rank(reports),
        "low_information_limits": list(first.low_information_limits) if first is not None else None,
        # checked once, on the first chunk, which holds the reads the CLI checks
        "strand": list(first.strand) if first is not None and first.strand else None,
    }


def withheld_by_rank(reports: Sequence[RunReport]) -> Dict[str, Dict[str, int]]:
    """Per rank and reason (``package_adapter.WITHHELD_REASONS``), the reads whose
    confidence reached the rank's threshold and that are Unclassified there,
    summed over the chunks of a job. Every such read is counted once per rank."""
    out = {lvl: {k: 0 for k in package_adapter.WITHHELD_REASONS} for lvl in LEVEL_ORDER}
    for r in reports:
        for lvl, counts in r.withheld_by_rank.items():
            for k, v in counts.items():
                out[lvl][k] += v
    return out