Blame for app.py

"""
Translatome -> microprotein -> candidate epitope, end to end.

A demo of what can be assembled *today* from existing, openly-licensed tools:

    transcript sequence
      -> TIS Transformer   P(translation initiation) per nucleotide   [MIT]
      -> ORF typing        uORF / intORF / N-ext / CDS ...            [here]
      -> MHCflurry         peptide presentation by HLA class I        [Apache-2.0]

Nothing in this Space is a new model. The point is the join, and what the join
reveals: how few candidates survive every stage, and which of the surviving
ones are actually novel rather than self.

Written as a starting point for a collaboration on beta-cell translatome
discovery -- hence the insulin-centric examples and the T1D risk alleles.
"""

import os
import tempfile
from pathlib import Path

import gradio as gr
import pandas as pd

import mhc
import pipeline

# --- Config -----------------------------------------------------------------

MAX_SEQS = 20
MAX_LEN = 30_000
MAX_TOTAL_LEN = 150_000
GPU_DURATION = 120
TIMEOUT_S = 600

# Weights ship inside the transcript-transformer wheel (~128 MB), so there is
# no download at first use and nothing to warm up.
ACCELERATOR = os.environ.get("TRISTAN_ACCELERATOR", "gpu")

EXAMPLE_FASTA = Path(__file__).parent / "examples.fa"

# Unlike transcript-transformer, MHCflurry does not ship its weights in the
# wheel. Fetch them once at container start rather than inside a request,
# where the download would eat a GPU allocation and time out. Set
# MHCFLURRY_BOOTSTRAP=0 locally if you have already fetched them.
if os.environ.get("MHCFLURRY_BOOTSTRAP", "1") == "1":
    ok, _ = mhc.predictor_available()
    if not ok:
        import subprocess

        print("Fetching MHCflurry models (one time, ~100s of MB)...", flush=True)
        subprocess.run(
            ["mhcflurry-downloads", "fetch", "models_class1_presentation"],
            check=False,
        )

try:
    import spaces

    HAS_ZEROGPU = True
except ImportError:
    HAS_ZEROGPU = False

    class _Stub:
        @staticmethod
        def GPU(*args, **kwargs):
            def wrap(fn):
                return fn

            return wrap

    spaces = _Stub()  # type: ignore


# --- The run ----------------------------------------------------------------


def _check_limits(records):
    if len(records) > MAX_SEQS:
        raise gr.Error(f"Too many sequences ({len(records)}). Limit is {MAX_SEQS}.")
    for name, seq in records:
        if len(seq) > MAX_LEN:
            raise gr.Error(
                f"Sequence '{name}' is {len(seq):,} nt. Limit is {MAX_LEN:,} nt."
            )
    total = sum(len(s) for _, s in records)
    if total > MAX_TOTAL_LEN:
        raise gr.Error(
            f"Total length {total:,} nt exceeds the {MAX_TOTAL_LEN:,} nt limit."
        )


@spaces.GPU(duration=GPU_DURATION)
def _tis_on_gpu(fasta_path, workdir, model, fold, cutoff, codons, min_len):
    # In-process: under ZeroGPU a child process may not inherit the
    # allocation, and would silently fall back to CPU.
    return pipeline.run_tis_transformer(
        Path(fasta_path), Path(workdir), model=model, fold=int(fold),
        accelerator="gpu", prob_cutoff=float(cutoff), start_codons=codons,
        min_orf_len=int(min_len), in_process=True,
    )


def _tis_on_cpu(fasta_path, workdir, model, fold, cutoff, codons, min_len):
    return pipeline.run_tis_transformer(
        Path(fasta_path), Path(workdir), model=model, fold=int(fold),
        accelerator="cpu", prob_cutoff=float(cutoff), start_codons=codons,
        min_orf_len=int(min_len), in_process=False, timeout_s=TIMEOUT_S,
    )


def run(fasta_text, uploaded, model, fold, cutoff, codon_mode, min_len,
        alleles, novel_only):
    if uploaded:
        fasta_text = Path(uploaded).read_text()

    try:
        records, headers = pipeline.parse_fasta_with_headers(fasta_text or "")
    except ValueError as e:
        raise gr.Error(str(e))
    _check_limits(records)

    if not alleles:
        raise gr.Error("Pick at least one HLA allele.")

    codons = (
        pipeline.ATG_ONLY if codon_mode.startswith("ATG")
        else pipeline.DEFAULT_START_CODONS
    )

    workdir = Path(tempfile.mkdtemp(prefix="tis_"))
    fasta_path = workdir / "input.fa"
    pipeline.write_fasta(records, fasta_path)

    # Score down to the floor whatever the user asked for, then apply their
    # cutoff here. Same work for the model, and it lets the run report what it
    # rejected instead of returning a bare "nothing found".
    cutoff = float(cutoff)
    floor = min(cutoff, pipeline.NEAR_MISS_FLOOR)

    # Try GPU first where there is one, then fall back to CPU. The model is
    # small enough that CPU takes seconds on inputs this size, so a GPU
    # problem should degrade the Space rather than break it.
    attempts = [_tis_on_gpu, _tis_on_cpu] if ACCELERATOR == "gpu" else [_tis_on_cpu]
    tis, failures = None, []
    for attempt in attempts:
        try:
            tis = attempt(str(fasta_path), str(workdir), model, fold, floor,
                          codons, min_len)
            break
        except Exception as e:
            failures.append(
                f"{attempt.__name__}: {type(e).__name__}: {e}".strip()
            )
            # Clear anything a half-finished run left behind.
            for stale in workdir.glob("tis.*"):
                stale.unlink(missing_ok=True)

    if tis is None:
        raise gr.Error(
            "Initiation-site prediction failed.\n\n" + "\n\n".join(failures)
        )

    total_nt = sum(len(s) for _, s in records)
    stages = [
        pipeline.Stage("Transcripts submitted", len(records), f"{total_nt:,} nt"),
    ]

    scored = pipeline.annotate_orf_types(tis, headers)
    ann, below = pipeline.split_at_cutoff(scored, cutoff)
    below_view = _tis_view(below).head(50)

    if ann.empty:
        stages.append(
            pipeline.Stage("Initiation sites", 0, f"none scored >= {cutoff}")
        )
        return (
            _funnel_md(stages, cutoff, alleles, _near_miss_note(below, cutoff)),
            pd.DataFrame(), below_view, pd.DataFrame(), None, None,
        )

    n_noncanon = int(ann["is_noncanonical"].sum())
    n_micro = int(ann["is_microprotein"].sum())

    stages += [
        pipeline.Stage("Initiation sites", len(ann), f"score >= {cutoff}"),
        pipeline.Stage("Non-canonical", n_noncanon, "not the annotated CDS start"),
        pipeline.Stage("Microproteins", n_micro, "< 100 aa"),
    ]

    # --- Stage 3 ---
    proteins = {
        f"{r.transcript_id.split('_')[0]}@{r.tis_pos_0b}[{r.orf_type}]":
            str(r.protein_sequence or "")
        for r in ann.itertuples()
    }
    res = mhc.predict_presentation(proteins, alleles=list(alleles))

    if res.error:
        stages.append(pipeline.Stage("Peptides scored", 0, res.error))
        mhc_table = pd.DataFrame()
    else:
        canon = pipeline.canonical_proteome(records, headers)
        tbl = pipeline.flag_peptide_novelty(res.table, canon)

        n_strong = int((tbl["binder"] == "strong").sum())
        stages += [
            pipeline.Stage("Peptides scored", len(tbl),
                           f"8-11mers x {len(alleles)} allele(s)"),
            pipeline.Stage("Strong binders", n_strong,
                           f"presentation percentile <= "
                           f"{mhc.STRONG_BINDER_PERCENTILE}"),
        ]

        if "novel" in tbl.columns:
            n_novel = int(((tbl["binder"] == "strong") & tbl["novel"]).sum())
            stages.append(
                pipeline.Stage("...and not in a canonical CDS", n_novel,
                               "absent from the submitted annotated proteins")
            )

        if novel_only and "novel" in tbl.columns:
            tbl = tbl[tbl["novel"]]
        mhc_table = tbl[tbl["binder"].isin(["strong", "weak"])].head(500)

    tis_view = _tis_view(ann)

    tmp = Path(tempfile.gettempdir())
    tis_csv = tmp / f"tis_sites_{os.getpid()}.csv"
    mhc_csv = tmp / f"peptides_{os.getpid()}.csv"
    tis_view.to_csv(tis_csv, index=False)
    if not mhc_table.empty:
        mhc_table.to_csv(mhc_csv, index=False)

    return (
        _funnel_md(stages, cutoff, alleles, _near_miss_note(below, cutoff)),
        tis_view,
        below_view,
        mhc_table,
        str(tis_csv),
        str(mhc_csv) if not mhc_table.empty else None,
    )


def _tis_view(df):
    """The columns worth showing for a set of scored initiation sites."""
    if df.empty:
        return pd.DataFrame()
    return df[[
        "transcript_id", "tis_pos_0b", "output", "start_codon", "orf_type",
        "protein_length", "TTS_on_transcript", "protein_sequence",
    ]].rename(columns={
        "tis_pos_0b": "TIS_pos_0based",
        "output": "TIS_score",
        "TTS_on_transcript": "has_stop",
    })


def _near_miss_note(below, cutoff):
    """One line on the best site the cutoff rejected.

    Without this, a transcript with no non-canonical hits looks like a tool
    that failed to look. It looked: the scores were simply tiny, which is the
    model's own verdict and the most informative thing it has to say.
    """
    if below.empty:
        return (
            f"Nothing at all scored above {pipeline.NEAR_MISS_FLOOR} outside "
            "what is shown — not a near miss in sight."
        )
    nc = below[below["is_noncanonical"]] if "is_noncanonical" in below else below
    if nc.empty:
        return ""
    top = nc.iloc[0]
    return (
        f"The best non-canonical site **below** the {cutoff} cutoff scored "
        f"**{top['output']:.4g}** ({top['orf_type']}, {top['start_codon']} at "
        f"position {int(top['tis_pos_0b'])} of `{top['transcript_id']}`). "
        "See the *Below the cutoff* tab — the size of that gap is the result."
    )


# Each example is (label, genes, fold). The fold is the one whose training and
# validation sets exclude that gene's chromosome -- scoring a transcript on a
# fold that trained on it inflates the result. Bundling the fold with the
# sequence means nobody has to know that to get a correct run.
#
# The order is deliberate. MIEF1 shows the model finding a validated 70 aa
# microprotein and scoring it above the canonical start, so the tool gets
# credit before any limitation is shown. ATF4 then exposes the limit inside a
# single transcript: uORF2 (59 aa) is found, uORF1 (3 aa) is not. TYK2 shows
# that limit landing on a type 1 diabetes risk gene. INS and IAPP are the
# negative controls, and the ones beta-cell people reach for first: on both,
# the annotated start is the only site worth anything, and the Below-the-cutoff
# tab shows by how much.
EXAMPLE_SETS = [
    ("Model finds a hidden microprotein, and misses a shorter one "
     "(ATF4 + MIEF1, chr22)", ["MIEF1", "ATF4"], 0),
    ("A validated upstream ORF the model scores at ~0 "
     "(TYK2, chr19, T1D risk gene)", ["TYK2"], 3),
    ("Reference: insulin, canonical start only (INS, chr11)", ["INS"], 4),
    ("Nothing non-canonical anywhere, and what that looks like "
     "(IAPP, chr12)", ["IAPP"], 3),
]


def example_rows():
    """Build (fasta_text, fold) rows and their labels from examples.fa."""
    if not EXAMPLE_FASTA.exists():
        return [], []

    blocks: dict[str, str] = {}
    for chunk in EXAMPLE_FASTA.read_text().split(">"):
        if not chunk.strip():
            continue
        gene = chunk.split()[0].split("_")[0]
        blocks[gene] = ">" + chunk.rstrip() + "\n"

    rows, labels = [], []
    for label, genes, fold in EXAMPLE_SETS:
        fasta = "".join(blocks[g] for g in genes if g in blocks)
        if fasta:
            rows.append([fasta, fold])
            labels.append(label)
    return rows, labels


def _funnel_md(stages, cutoff, alleles, near_miss=""):
    """Render the attrition, because the attrition is the honest result."""
    lines = [
        "### How many candidates survive each stage",
        "",
        "| stage | n | filter |",
        "| --- | ---: | --- |",
    ]
    for s in stages:
        lines.append(f"| {s.name} | **{s.count:,}** | {s.detail} |")
    lines += [
        "",
        f"*Alleles: {', '.join(alleles)}. TIS cutoff {cutoff}.*",
        "",
        "Each arrow multiplies one model's error by another's. Treat the final "
        "row as a shortlist to test, not as predicted neoantigens.",
    ]
    if near_miss:
        lines += ["", near_miss]
    return "\n".join(lines)


# --- UI ---------------------------------------------------------------------

with gr.Blocks(title="Translatome to epitope") as demo:
    gr.Markdown(
        """
        # Translatome → microprotein → candidate epitope

        What you can already build today, end to end, from open tools:

        **transcript sequence** → `TIS Transformer` (predicted initiation sites,
        MIT) → **ORF typing** → `MHCflurry` (HLA class I presentation,
        Apache-2.0) → **candidate peptides**

        No new models here. The interesting part is what the join shows: how
        steeply candidates drop off at each stage, and how many survivors are
        merely *self* peptides rather than novel ones.
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            fasta_text = gr.Textbox(
                label="Transcript sequence(s), FASTA",
                placeholder=">my_transcript 5UTR=59nt CDS=333nt\nACGT...",
                info="Tag a header with 5UTR=<n>nt CDS=<n>nt to get ORF types "
                     "(uORF, intORF, N-ext). Without tags, ORFs are reported "
                     "as unannotated.",
                lines=10,
            )
            uploaded = gr.File(
                label="or upload a .fa file",
                file_types=[".fa", ".fasta", ".txt"],
            )
        with gr.Column(scale=2):
            with gr.Accordion("Initiation site prediction", open=True):
                model = gr.Radio(["human", "mouse"], value="human", label="Model")
                fold = gr.Dropdown(
                    [0, 1, 2, 3, 4], value=0, label="Fold",
                    info="Each fold held out different chromosomes during "
                         "training. Pick one that did NOT train on your "
                         "sequence's chromosome, or the score is inflated. "
                         "For the examples: TYK2 and IAPP → 3, INS → 4. "
                         "Fold 0 is the catch-all and will score anything.",
                )
                cutoff = gr.Slider(
                    0.0005, 1.0, value=pipeline.DEFAULT_PROB_CUTOFF, step=0.0005,
                    label="TIS score cutoff",
                    info=f"Package default 0.03. Scoring always runs down to "
                         f"{pipeline.NEAR_MISS_FLOOR}; everything this cutoff "
                         f"rejects is listed in the Below-the-cutoff tab.",
                )
                codon_mode = gr.Radio(
                    ["Near-cognate (ATG, CTG, GTG, TTG)", "ATG only"],
                    value="Near-cognate (ATG, CTG, GTG, TTG)",
                    label="Start codons",
                    info="Near-cognate is the package default and is where "
                         "much non-canonical initiation happens.",
                )
                min_len = gr.Slider(9, 300, value=pipeline.DEFAULT_MIN_ORF_LEN,
                                    step=3, label="Min ORF length (nt)")
            with gr.Accordion("Epitope prediction", open=True):
                alleles = gr.CheckboxGroup(
                    choices=list(mhc.ALL_ALLELES.keys()),
                    value=mhc.DEFAULT_ALLELES,
                    label="HLA class I alleles",
                    info="The first four are T1D-associated.",
                )
                novel_only = gr.Checkbox(
                    value=False,
                    label="Show only peptides absent from the canonical CDS",
                    info="Peptides shared with the canonical protein are self "
                         "and are not neoantigen candidates.",
                )
            go = gr.Button("Run pipeline", variant="primary")

    funnel = gr.Markdown("")

    with gr.Tab("Initiation sites"):
        tis_out = gr.Dataframe(label="Predicted initiation sites", wrap=True)
        tis_dl = gr.File(label="Download (CSV)")
    with gr.Tab("Below the cutoff"):
        gr.Markdown(
            f"Every run scores down to {pipeline.NEAR_MISS_FLOOR} regardless "
            "of the cutoff, so these are the sites the cutoff rejected, best "
            "first. On most transcripts the only non-canonical candidates "
            "live here, two or three orders of magnitude below the annotated "
            "start — that gap is the model telling you it sees nothing, not "
            "the pipeline failing to look."
        )
        below_out = gr.Dataframe(label="Rejected sites (top 50)", wrap=True)
    with gr.Tab("Candidate peptides"):
        mhc_out = gr.Dataframe(label="Strong and weak binders", wrap=True)
        mhc_dl = gr.File(label="Download (CSV)")
    with gr.Tab("What this does not show"):
        gr.Markdown(
            """
            Worth being explicit, because a chained pipeline flatters itself.

            **Sequence alone cannot be cell-type specific.** TIS Transformer
            sees only sequence, so it returns the same answer for a beta cell,
            a hepatocyte and a neuron. Beta-cell-specific initiation under
            inflammatory stress — the thing that actually matters for type 1
            diabetes — is invisible to it. That needs ribosome profiling or
            expression conditioning.

            **It is largely a canonical-start detector.** Trained on annotated
            CDS starts, it reproduces them confidently and assigns most
            non-canonical sites a score near zero. On insulin, scored right
            down to 0.0005 on the fold that held out chr11, exactly two sites
            exist in the whole transcript: the annotated ATG at 0.91 and an
            in-frame N-truncation at 0.0037. Nothing in another frame scores
            at all. IAPP behaves the same way — annotated start 0.18, then one
            3'UTR ORF at 0.0008.

            That is worth separating from the ribosome-profiling point above.
            A blank result on a gene with a documented alternative-frame
            product, such as the insulin DRiP of
            [Kracht et al.](https://doi.org/10.1038/nm.4289), is not caused by
            the missing Ribo-seq track: the model was shown the sequence and
            returned a near-zero probability. Conditioning on ribosome
            occupancy might change that answer, but the limit you are looking
            at here is the model's prior, not a missing input.

            **Novelty here is local.** The "novel" flag compares a peptide only
            against the canonical CDSs of the sequences you submitted. A real
            screen needs the whole human proteome, and ideally the canonical
            immunopeptidome too.

            **Class I only.** MHCflurry does not model class II, and the
            strongest genetic risk in T1D is class II (HLA-DQ8/DQ2).

            **Binding is not immunogenicity.** Presentation prediction says a
            peptide could be displayed. Whether a T cell responds depends on
            the repertoire, tolerance and much else.

            **No proteomic evidence is used.** Peptide-level support for
            non-canonical ORFs now exists at scale (the TransCODE consensus
            set); none of it is consulted here.
            """
        )

    gr.Markdown(
        f"""
        ---
        **Limits:** {MAX_SEQS} sequences, {MAX_LEN:,} nt each,
        {MAX_TOTAL_LEN:,} nt total.

        **Not included:** RiboTIE, the ribosome-profiling model from the same
        package. It needs mapped BAMs that are too large for a Space.

        TIS Transformer: Clauwaert et al., *NAR Genomics and Bioinformatics*
        (2023), [doi:10.1093/nargab/lqad021](https://doi.org/10.1093/nargab/lqad021).
        MHCflurry: O'Donnell et al., *Cell Systems* (2020),
        [doi:10.1016/j.cels.2020.06.010](https://doi.org/10.1016/j.cels.2020.06.010).
        """
    )

    go.click(
        run,
        inputs=[fasta_text, uploaded, model, fold, cutoff, codon_mode, min_len,
                alleles, novel_only],
        outputs=[funnel, tis_out, below_out, mhc_out, tis_dl, mhc_dl],
    )

    rows, labels = example_rows()
    if rows:
        try:
            gr.Examples(
                examples=rows,
                inputs=[fasta_text, fold],
                example_labels=labels,
                label="Examples (each preselects the fold that held out its "
                      "chromosome)",
            )
        except TypeError:
            # example_labels landed in Gradio 4.x; degrade rather than break.
            gr.Examples(
                examples=rows, inputs=[fasta_text, fold],
                label="Examples (each preselects its correct fold)",
            )


if __name__ == "__main__":
    demo.queue(max_size=20).launch()