biolmnet/ui/components.py

"""HTML builders for the Workbench shell.

Every function here returns a plain HTML string for a ``gr.HTML`` component.
They only format data already produced by ``biolmnet.data`` /
``biolmnet.training`` / ``biolmnet.artifacts`` — none of them compute
anything scientific.
"""

from __future__ import annotations

import html
from typing import Iterable, Sequence

esc = html.escape

CORNER_MARKS = (
    '<i class="corner tl"></i><i class="corner tr"></i>'
    '<i class="corner bl"></i><i class="corner br"></i>'
)


def blueprint_div(inner_html: str, *, extra_class: str = "", style: str = "") -> str:
    """Wrap ``inner_html`` in a blueprint frame (hairline border + corner marks)."""
    cls = f"blueprint {extra_class}".strip()
    style_attr = f' style="{style}"' if style else ""
    return f'<div class="{cls}"{style_attr}>{CORNER_MARKS}{inner_html}</div>'


# ── rail ─────────────────────────────────────────────────────────────────

def brand_block_html() -> str:
    return (
        '<div class="rail-brand"><span class="name">BioLM-NET</span>'
        '<div class="mono" style="margin-top:5px">Workbench</div></div>'
        '<div class="mono rail-kicker">Workflow</div>'
    )


def rail_row_html(number: int, label: str, state: str, status_word: str) -> str:
    """One workflow rail row. ``state`` is one of done/on/next/off."""
    return (
        f'<div class="rr {state}">'
        f'<span class="rn">{number:02d}</span>'
        f'<span class="lb">{esc(label)}</span>'
        f'<span class="st">{esc(status_word)}</span>'
        f"</div>"
    )


def kv_row(key: str, value: str, tone: str | None = None) -> str:
    cls = f"num {tone}" if tone else "num"
    return f'<div class="kv"><span>{esc(key)}</span><span class="{cls}">{esc(value)}</span></div>'


def run_state_plate(rows: Sequence[tuple[str, str, str | None]]) -> str:
    """``rows``: (key, value, tone) where tone is None/accent/error/muted."""
    inner = "".join(kv_row(k, v, tone) for k, v, tone in rows)
    return blueprint_div(inner, extra_class="plate")


def kv_plain_row(key: str, value: str, *, last: bool = False) -> str:
    cls = "kv-plain last" if last else "kv-plain"
    return f'<div class="{cls}"><span>{esc(key)}</span><span class="num">{esc(value)}</span></div>'


def footnote_html(line1: str, line2: str) -> str:
    return f'<div class="rail-footnote">{esc(line1)}<br>{esc(line2)}</div>'


# ── topbar / page head ──────────────────────────────────────────────────

def topbar_html(stage_no: int, total: int, stage_name: str, session_id: str, right_text: str) -> str:
    # Only the left "STAGE 0n / 05" label is uppercase; the session/device
    # note on the right stays mixed-case (a lowercase-hex session id reading
    # upper-cased is just noise).
    return (
        '<div class="topbar">'
        f'<div class="mono">Stage {stage_no:02d} / {total:02d} — {esc(stage_name)}</div>'
        f'<div class="num" style="font-size:9.5px;letter-spacing:.04em;color:var(--color-muted-2)">'
        f"Session {esc(session_id)} · {esc(right_text)}</div>"
        "</div>"
    )


def title_block_html(title: str, description: str) -> str:
    return f"<h2>{esc(title)}</h2>" f'<div class="desc">{esc(description)}</div>'


def mono_meta_html(lines: Iterable[str], *, align_right: bool = True) -> str:
    align = "text-align:right;" if align_right else ""
    body = "<br>".join(esc(line) for line in lines)
    return f'<div class="mono" style="{align}font-size:9.5px;line-height:1.7">{body}</div>'


# ── stat plates ──────────────────────────────────────────────────────────

def stat_plate(value: str, label: str, *, label_first: bool = False) -> str:
    if label_first:
        inner = f'<div class="k">{esc(label)}</div><div class="v">{esc(value)}</div>'
    else:
        inner = f'<div class="v">{esc(value)}</div><div class="k">{esc(label)}</div>'
    return blueprint_div(inner, extra_class="stat")


def mini_stat_html(value: str, label: str) -> str:
    """A label+value pair with no frame of its own (used inside a shared blueprint panel)."""
    return (
        f'<div><div class="mono" style="font-size:9px">{esc(label)}</div>'
        f'<div style="font:600 26px/1.1 var(--font-heading)">{esc(value)}</div></div>'
    )


# ── bars ─────────────────────────────────────────────────────────────────

def bar_cell(pct: float, *, height: int = 8) -> str:
    pct = max(0.0, min(100.0, pct))
    return f'<div class="bar" style="height:{height}px"><i style="width:{pct:.1f}%"></i></div>'


def labeled_bar_row(
    label: str, pct: float, value_text: str, *, label_width: int = 78, bar_height: int = 12, value_width: int = 52
) -> str:
    return (
        '<div style="display:flex;align-items:center;gap:10px">'
        f'<span style="width:{label_width}px;font-size:12px">{esc(label)}</span>'
        f"{bar_cell(pct, height=bar_height)}"
        f'<span class="num" style="width:{value_width}px;text-align:right">{esc(value_text)}</span>'
        "</div>"
    )


# ── tables ───────────────────────────────────────────────────────────────

def table_html(
    headers: Sequence[str],
    rows: Sequence[Sequence[str]],
    *,
    aligns: Sequence[str] | None = None,
    numeric_cols: Sequence[bool] | None = None,
    widths: Sequence[str | None] | None = None,
) -> str:
    """Build a ``.tbl`` table. Cell values are raw HTML — escape plain text
    with :data:`esc` before passing it in."""
    n_cols = len(headers)
    aligns = list(aligns) if aligns else ["left"] * n_cols
    numeric_cols = list(numeric_cols) if numeric_cols else [False] * n_cols
    widths = list(widths) if widths else [None] * n_cols

    head_cells = []
    for header, align, width in zip(headers, aligns, widths):
        style_parts = []
        if width:
            style_parts.append(f"width:{width}")
        if align == "right":
            style_parts.append("text-align:right")
        attr = f' style="{";".join(style_parts)}"' if style_parts else ""
        head_cells.append(f"<th{attr}>{esc(header)}</th>")

    body_rows = []
    for row in rows:
        cells = []
        for value, align, numeric in zip(row, aligns, numeric_cols):
            cls_attr = ' class="n"' if numeric else ""
            style_attr = ' style="text-align:right"' if align == "right" else ""
            cells.append(f"<td{cls_attr}{style_attr}>{value}</td>")
        body_rows.append(f"<tr>{''.join(cells)}</tr>")

    return (
        '<table class="tbl"><thead><tr>'
        + "".join(head_cells)
        + "</tr></thead><tbody>"
        + "".join(body_rows)
        + "</tbody></table>"
    )


def confusion_matrix_html(matrix, labels: Sequence[str]) -> str:
    """A tinted CSS grid confusion matrix (rows true, columns predicted)."""
    n = len(labels)
    row_totals = [max(sum(row), 1) for row in matrix]

    def cell_style(row_index: int, count: float) -> str:
        ratio = count / row_totals[row_index]
        if ratio >= 0.7:
            return "background:var(--color-accent-800);color:#f2f2f3"
        if ratio >= 0.4:
            return "background:var(--color-accent-700);color:#f2f2f3"
        if ratio >= 0.15:
            return "background:var(--color-accent-200)"
        return "background:var(--color-accent-100)"

    header_cells = "".join(
        f'<div class="mono" style="font-size:8.5px;text-align:center">{esc(label)}</div>' for label in labels
    )
    row_label_cells = "".join(
        f'<div class="mono" style="font-size:8.5px;display:grid;align-items:center">{esc(label)}</div>'
        for label in labels
    )
    body_cells = "".join(
        f'<div style="{cell_style(i, matrix[i][j])}">{int(matrix[i][j])}</div>'
        for i in range(n)
        for j in range(n)
    )
    return (
        '<div style="display:grid;grid-template-columns:56px 1fr;gap:8px">'
        "<div></div>"
        f'<div style="display:grid;grid-template-columns:repeat({n},1fr);gap:2px">{header_cells}</div>'
        f'<div style="display:grid;gap:2px">{row_label_cells}</div>'
        f'<div class="cm" style="--cm-n:{n}">{body_cells}</div>'
        "</div>"
    )


# ── status / validation strips ───────────────────────────────────────────

def simple_status_html(message_html: str, *, error: bool = False) -> str:
    cls = "error" if error else "status"
    return f'<div class="strip {cls}"><div class="strip-body" style="flex:1">{message_html}</div></div>'


def strip_text_html(kicker: str, message: str) -> str:
    """The text half of a bordered strip; pair with a real ``gr.Button`` in the
    same ``elem_classes=["strip", ...]`` row for the ghost action."""
    return f'<div class="strip-kicker">{esc(kicker)}</div><div class="strip-body">{esc(message)}</div>'


def empty_note_html(message: str) -> str:
    return f'<div class="empty-note">{esc(message)}</div>'


# ── introduction cards ────────────────────────────────────────────────────

INTRO_CARDS = [
    (
        "Introduction · 01 / 06",
        "A network wired along known biology",
        (
            "BioLM-NET classifies tumour samples from <b>paired gene expression and DNA methylation</b>. "
            "Instead of connecting every gene to every neuron, it only makes the connections biology supports — "
            "KEGG pathways, transcription-factor targets, protein interactions — so trained weights can be read "
            "back as pathways rather than as a black box."
        ),
        "This workbench runs the model end to end in five stages. On the bundled BRCA example it takes a few minutes.",
        "",
    ),
    (
        "Introduction · 02 / 06",
        "Five stages, in order",
        (
            "The rail on the left is the spine of the app. Each stage consumes the previous one's output, so a "
            "stage stays <b>locked</b> until its prerequisite exists — you cannot export a model you have not trained."
        ),
        (
            '<div class="intro-stage-list">'
            "<div><em>01</em><span><b>Data &amp; Priors</b> — assemble the masked graph</span></div>"
            "<div><em>02</em><span><b>Train Model</b> — fit it, watch the epochs</span></div>"
            "<div><em>03</em><span><b>Export Artifacts</b> — save a reusable bundle</span></div>"
            "<div><em>04</em><span><b>Predict</b> — score a new cohort</span></div>"
            "<div><em>05</em><span><b>Results</b> — metrics and pathway attention</span></div>"
            "</div>"
        ),
        "The plate under the rail always says what state the model is in",
    ),
    (
        "Introduction · 03 / 06 · Stage 01",
        "Data & Priors",
        (
            "Point the workbench at paired omics — a bundled example, a GitHub folder, or your own upload — with "
            "<b>samples in rows and HGNC symbols in columns</b>. The two matrices are aligned on their shared "
            "samples, and every file is reported with the shape it was read at rather than failing silently."
        ),
        (
            "The priors on the right are the biology that becomes wiring: pathway annotations, the enrichment "
            "cutoff, the GenePT context and the interaction sources. Defaults follow the paper."
        ),
        "Rebuilding clears any trained model",
    ),
    (
        "Introduction · 04 / 06 · Stage 02",
        "Train Model",
        (
            "Hyperparameters sit on the left as a spec sheet. The run panel on the right reports <b>epoch, train "
            "and validation loss, accuracy and ETA</b> as the fit progresses."
        ),
        (
            "Controls lock for the duration of a run and the primary button reads <b>Training on ZeroGPU</b> "
            "until it finishes; Cancel stays available. The GPU is allocated on demand, so the first epoch may "
            "wait for its reservation."
        ),
        "Paper defaults are one click away",
    ),
    (
        "Introduction · 05 / 06 · Stages 03–04",
        "Export, then predict",
        (
            "Export writes one bundle carrying the weights, the fitted preprocessing, the mask and the config — "
            "enough to score new samples later without rebuilding the graph. Its manifest lists what each entry "
            "reproduces, with a checksum."
        ),
        (
            "Predict scores a new cohort with either the session model or an uploaded bundle. Features are "
            "<b>aligned to the artifact first</b>; if a required column is missing, inference is blocked and the "
            "check that failed is named."
        ),
        "Bundles are temp files — download before the Space restarts",
    ),
    (
        "Introduction · 06 / 06 · Stage 05",
        "Reading the results",
        (
            "Results collects everything a run produced: validation metrics, the training curve, the confusion "
            "matrix, an audit of how sparse each layer actually is, and the pathway attention weights — which is "
            "the part the architecture exists to give you."
        ),
        (
            "Treat it as model output, not as biology. <b>Check cohort composition, preprocessing and class "
            "balance before drawing conclusions</b>, and remember that attention weights rank pathways within "
            "this fit rather than proving mechanism."
        ),
        "Research use only · not for clinical decisions",
    ),
]


def intro_card_html(index: int) -> str:
    step, title, body_1, body_2, _ = INTRO_CARDS[index]
    return (
        f'<div class="intro-step">{esc(step)}</div>'
        f"<h4>{esc(title)}</h4>"
        f"<p>{body_1}</p>"
        f"<p>{body_2}</p>"
    )


def intro_dots_html(index: int, total: int = 6) -> str:
    return '<div class="intro-dots">' + "".join(
        f'<i class="{"on" if i == index else ""}"></i>' for i in range(total)
    ) + "</div>"


# ── misc ─────────────────────────────────────────────────────────────────

def human_bytes(n: float) -> str:
    n = float(n)
    for unit in ("B", "KB", "MB"):
        if n < 1024 or unit == "MB":
            return f"{int(n)} {unit}" if unit == "B" else f"{n:.1f} {unit}"
        n /= 1024
    return f"{n / 1024:.1f} GB"