src/uploads/assembly.py

"""Assemble a counts/TPM/metadata delivery into an analysis-ready h5ad.

This is the library form of what `scripts/assemble_myc_kd_kmc_mouse.py` used to
do entirely inline: turn a Novogene-style delivery

    gene_level_counts.tsv       genes x samples, ENSMUSG ids, R write.table
                                header (first column name missing)
    gene_level_abundances.tsv   same shape, TPM (optional; stored as a layer)
    sample metadata             .xlsx / .csv / .tsv, must yield the obs columns
                                clone / arm / site / mouse_id

into

    X               int32 rounded counts, samples x genes (Path A / DESeq2)
    layers['tpm']   float32 TPM
    var.index       MGI mouse symbol (duplicates summed)
    obs             clone, arm, site, mouse_id (categoricals) + extras
    uns             organism='mouse', analysis_space='mouse'

Nothing here calls ``sys.exit`` — every input problem raises
:class:`AssemblyError`, so the same code can back a CLI *and* the ADR-0011
upload path (where a bad sheet must become a message in the UI, not a dead
process). The CLI wrapper converts the exception back into an exit.

Safety note: these readers are the vetted-loader equivalent for the upload gate
— pandas/anndata parsing only. Nothing is ``exec``'d, and the caller is expected
to hand over paths that have already been staged and content-scanned.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

REPO_ROOT = Path(__file__).resolve().parent.parent.parent
SYMBOL_MAP_PATH = REPO_ROOT / "resources" / "mouse_ensembl_symbol_map.tsv.gz"
REQUIRED_OBS = ["clone", "arm", "site", "mouse_id"]

EXCEL_SUFFIXES = (".xlsx", ".xls")


class AssemblyError(Exception):
    """A user-fixable problem with the supplied files (bad column, no join)."""


def load_mouse_symbol_map(path: Path = SYMBOL_MAP_PATH) -> pd.Series:
    """ENSMUSG (unversioned) -> MGI symbol, as a Series."""
    df = pd.read_csv(path, sep="\t", comment="#")
    return pd.Series(df["mouse_symbol"].values, index=df["ensembl_gene_id"].values)


def map_mouse_ensembl_to_symbols(matrix: pd.DataFrame) -> tuple[pd.DataFrame, dict]:
    """
    Map a genes x samples matrix from ENSMUSG ids to MGI symbols.

    Strips Ensembl version suffixes, drops unmapped genes, and SUMS rows that
    collapse to the same symbol (correct for counts; acceptable for TPM since
    multi-locus duplicates are a handful of rows). Returns (mapped matrix
    indexed by symbol with an 'ensembl_gene_id' representative kept in
    .attrs['ensembl_of_symbol'], stats dict).
    """
    symap = load_mouse_symbol_map()
    stripped = matrix.index.astype(str).str.split(".").str[0]
    symbols = stripped.map(symap)  # Index of symbols (NaN where unmapped)
    mask = pd.notna(symbols)
    mapped = matrix.loc[mask].copy()
    mapped.index = symbols[mask]
    n_dup = int(mapped.index.duplicated().sum())
    ensembl_of_symbol = pd.Series(stripped[mask], index=mapped.index).groupby(level=0).first()
    collapsed = mapped.groupby(level=0).sum()
    collapsed.attrs["ensembl_of_symbol"] = ensembl_of_symbol
    stats = {
        "n_input_genes": int(matrix.shape[0]),
        "n_unmapped_dropped": int(symbols.isna().sum()),
        "n_duplicate_rows_summed": n_dup,
        "n_output_symbols": int(collapsed.shape[0]),
    }
    return collapsed, stats


def parse_column_map(spec: str | None) -> dict[str, str]:
    """'mouse_id=Mouse,site=Type' -> {target obs col: source sheet col}."""
    if not spec:
        return {}
    out: dict[str, str] = {}
    for pair in spec.split(","):
        if not pair.strip():
            continue
        if "=" not in pair:
            raise AssemblyError(f"Bad column-map entry {pair!r}; expected target=Source.")
        target, source = pair.split("=", 1)
        out[target.strip()] = source.strip()
    return out


def parse_value_maps(specs: list[str] | str | None) -> dict[str, dict[str, str]]:
    """['site=Tumor:tumor', 'site=Met:liver_met'] -> {'site': {'Tumor': 'tumor', ...}}.

    A single string is accepted too (newline- or semicolon-separated), which is
    what a one-line UI text box produces.
    """
    if isinstance(specs, str):
        specs = [s for s in specs.replace(";", "\n").splitlines() if s.strip()]
    out: dict[str, dict[str, str]] = {}
    for spec in specs or []:
        if "=" not in spec or ":" not in spec.split("=", 1)[1]:
            raise AssemblyError(f"Bad value-map entry {spec!r}; expected column=old:new.")
        col, mapping = spec.split("=", 1)
        old, new = mapping.split(":", 1)
        out.setdefault(col.strip(), {})[old.strip()] = new.strip()
    return out


def detect_header_row(raw: pd.DataFrame, max_scan: int = 20) -> int:
    """First row (within max_scan) with no empty cells — Novogene sheets carry
    a short free-text preamble above the real header. Falls back to 0."""
    for i in range(min(max_scan, len(raw))):
        row = raw.iloc[i]
        if row.notna().all() and not row.astype(str).str.strip().eq("").any():
            return i
    return 0


def load_metadata(
    path: Path,
    sample_column: str | None,
    skip_rows: int | None = None,
    column_map: dict[str, str] | None = None,
    value_maps: dict[str, dict[str, str]] | None = None,
    group_column: str | None = None,
    control_label: str | None = None,
    treatment_label: str = "shMyc",
) -> pd.DataFrame:
    path = Path(path)
    is_excel = path.suffix.lower() in EXCEL_SUFFIXES
    if is_excel:
        try:
            if skip_rows is None:
                raw = pd.read_excel(path, header=None)
                skip_rows = detect_header_row(raw)
            meta = pd.read_excel(path, skiprows=skip_rows)
        except ImportError as e:
            raise AssemblyError(
                f"Reading {path.name} needs openpyxl ({e}). Export the sheet to CSV and retry."
            ) from e
    else:
        meta = pd.read_csv(
            path,
            sep="\t" if path.suffix.lower() in (".tsv", ".txt") else ",",
            skiprows=skip_rows or 0,
        )
    meta.columns = [str(c).strip() for c in meta.columns]

    # Rename mapped source columns to their target obs names (before the index
    # is set, so the sample column itself may be remapped too).
    for target, source in (column_map or {}).items():
        if source not in meta.columns:
            raise AssemblyError(
                f"Column-map source column {source!r} not in sheet. Present: {list(meta.columns)}."
            )
        meta[target] = meta[source]

    # Derive arm + clone from a single group column (control rows: clone=none).
    if group_column:
        if control_label is None:
            raise AssemblyError("A group column requires a control label.")
        if group_column not in meta.columns:
            raise AssemblyError(
                f"Group column {group_column!r} not in sheet. Present: {list(meta.columns)}."
            )
        group = meta[group_column].astype(str).str.strip()
        is_control = group == control_label
        if not is_control.any():
            raise AssemblyError(
                f"Control label {control_label!r} matches no rows of "
                f"{group_column!r} (values: {sorted(group.unique())})."
            )
        meta["arm"] = np.where(is_control, control_label, treatment_label)
        meta["clone"] = np.where(is_control, "none", group)

    id_col = sample_column or meta.columns[0]
    if id_col not in meta.columns:
        raise AssemblyError(
            f"Sample column {id_col!r} not in sheet. Present: {list(meta.columns)}."
        )
    meta = meta.set_index(meta[id_col].astype(str).str.strip()).drop(columns=[id_col])
    meta.columns = [c.strip().lower().replace(" ", "_") for c in meta.columns]

    for col, mapping in (value_maps or {}).items():
        if col not in meta.columns:
            raise AssemblyError(
                f"Value-map column {col!r} not in metadata. Present: {list(meta.columns)}."
            )
        meta[col] = meta[col].astype(str).str.strip().replace(mapping)

    missing = [c for c in REQUIRED_OBS if c not in meta.columns]
    if missing:
        raise AssemblyError(
            f"Metadata is missing required column(s) {missing}. Present: "
            f"{list(meta.columns)}. Rename them in the sheet, or use the column-map "
            f"/ group-column options."
        )
    return meta


def assemble_h5ad(
    counts_path: str | Path,
    metadata_path: str | Path,
    *,
    tpm_path: str | Path | None = None,
    out_path: str | Path | None = None,
    sample_column: str | None = None,
    skip_rows: int | None = None,
    column_map: str | dict[str, str] | None = None,
    value_maps: str | list[str] | dict[str, dict[str, str]] | None = None,
    group_column: str | None = None,
    control_label: str | None = None,
    treatment_label: str = "shMyc",
    staging_script: str = "src/uploads/assembly.py",
) -> tuple[Any, dict]:
    """Build the analysis h5ad from counts (+ optional TPM) and a metadata sheet.

    Returns ``(adata, report)``. ``report`` carries the per-matrix mapping stats,
    the obs value counts, and any unmatched sample ids — the same facts the CLI
    used to print, so a UI can show them instead.

    Writes to ``out_path`` when given; otherwise the AnnData is returned only.
    """
    import anndata as ad

    if isinstance(column_map, str) or column_map is None:
        column_map = parse_column_map(column_map)
    if not isinstance(value_maps, dict):
        value_maps = parse_value_maps(value_maps)

    counts = pd.read_csv(counts_path, sep="\t", index_col=0)  # absorbs the R header
    if counts.empty:
        raise AssemblyError(f"{Path(counts_path).name} has no data rows.")
    counts_sym, counts_stats = map_mouse_ensembl_to_symbols(counts)
    if counts_sym.empty:
        raise AssemblyError(
            f"No gene id in {Path(counts_path).name} mapped to a mouse symbol — "
            "the matrix does not look like unversioned/versioned ENSMUSG ids."
        )

    meta = load_metadata(
        Path(metadata_path),
        sample_column,
        skip_rows=skip_rows,
        column_map=column_map,
        value_maps=value_maps,
        group_column=group_column,
        control_label=control_label,
        treatment_label=treatment_label,
    )

    samples = [s for s in counts_sym.columns if s in meta.index]
    unmatched = [s for s in counts_sym.columns if s not in meta.index]
    if not samples:
        raise AssemblyError(
            "No matrix sample id matches the metadata index — check the sample "
            f"column. Matrix ids: {list(counts_sym.columns)[:5]}…; "
            f"metadata ids: {list(meta.index)[:5]}…"
        )

    X = counts_sym[samples].T
    adata = ad.AnnData(
        X=np.rint(X.values).astype(np.int32),
        obs=meta.loc[samples].copy(),
        var=pd.DataFrame(
            {"ensembl_gene_id": counts_sym.attrs["ensembl_of_symbol"].reindex(X.columns).values},
            index=pd.Index(X.columns, name="mouse_symbol"),
        ),
    )
    for col in REQUIRED_OBS:
        adata.obs[col] = adata.obs[col].astype(str).str.strip().astype("category")

    tpm_stats = None
    if tpm_path is not None:
        tpm = pd.read_csv(tpm_path, sep="\t", index_col=0)
        tpm_sym, tpm_stats = map_mouse_ensembl_to_symbols(tpm)
        adata.layers["tpm"] = (
            tpm_sym.reindex(index=X.columns, columns=samples)
            .fillna(0.0)
            .T.values.astype(np.float32)
        )

    adata.uns["organism"] = "mouse"
    adata.uns["analysis_space"] = "mouse"
    adata.uns["staging_script"] = staging_script

    report = {
        "counts_stats": counts_stats,
        "tpm_stats": tpm_stats,
        "n_samples": int(adata.n_obs),
        "n_genes": int(adata.n_vars),
        "unmatched_samples": unmatched,
        "obs_columns": list(adata.obs.columns),
        "obs_counts": {c: dict(adata.obs[c].value_counts()) for c in REQUIRED_OBS},
    }

    if out_path is not None:
        adata.write_h5ad(out_path)
        report["out_path"] = str(out_path)
    return adata, report