src/uploads/drafting.py

"""Draft a manifest from an uploaded file (ADR-0011 friction mitigation).

The gate requires a manifest (item 2), which is the main friction for a
non-coding uploader. This helper does the ADR's "drafts the manifest from the
file + a few prompts": it opens the file with the SAME vetted loaders the
validator uses (``scanpy.read_h5ad`` / ``pandas.read_csv`` — never
``exec``/``eval``/``pickle``), infers what it safely can (data level, feature-ID
type, sample/feature counts, candidate grouping columns), and returns a manifest
skeleton plus an explicit ``todo`` list of the fields a human must still confirm.

It does NOT weaken the gate: a drafted manifest is just a starting point that
still has to clear :func:`validate_upload` (which re-checks these same
inferences against the data) and admin registration. It only removes the
blank-page problem, so the honest workflow is:

    draft = draft_manifest(path)      # inspect + pre-fill
    # uploader reviews draft.todo, edits draft.manifest
    rec = stage_upload(path, manifest=draft.manifest, ...)
    validate_upload(rec)              # re-verifies the inferences

A backing UI (local Gradio panel) can wrap this to present ``manifest`` as an
editable form with ``todo`` as the required-fields checklist; the inference is
here so the UI stays a thin shell.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

# detected data_type (microarray classifier) → manifest data_level vocabulary.
_DATA_LEVEL_MAP = {
    "raw_counts": "raw_counts",
    "log_expression": "log_expression",
    "log_ratio_microarray": "log_ratio",
    "unknown": None,
}

# detected feature pattern (check_feature_id_type) → manifest feature_id_type.
_FEATURE_MAP = {
    "gene_symbol": "gene_symbol",
    "ensembl": "ensembl_gene_id",
    "entrez": "entrez_id",
    "probe_id_or_unknown": None,
}

_TABULAR_SUFFIXES = (".csv", ".tsv", ".txt", ".csv.gz", ".tsv.gz", ".txt.gz")
_MAX_X = 50_000
_TODO = "TODO"


@dataclass
class DraftResult:
    """A drafted manifest plus what still needs a human."""

    manifest: dict = field(default_factory=dict)
    inferred: dict = field(default_factory=dict)  # field → (value, evidence)
    todo: list[str] = field(default_factory=list)  # required fields still unfilled
    notes: list[str] = field(default_factory=list)


def _kind_of(name: str) -> str | None:
    low = name.lower()
    if low.endswith(".h5ad"):
        return "h5ad"
    if low.endswith(_TABULAR_SUFFIXES):
        return "tabular"
    return None


def _cap_flat(X_flat):
    import numpy as np

    if len(X_flat) > _MAX_X:
        return np.random.default_rng(0).choice(X_flat, size=_MAX_X, replace=False)
    return X_flat


def _infer_data_level(X_flat) -> tuple[str | None, str]:
    from src.workflows.microarray import classify_expression_data_type

    info = classify_expression_data_type(X_flat)
    detected = info["data_type"]
    level = _DATA_LEVEL_MAP.get(detected)
    ev = (
        f"detected '{detected}' (integer={info['is_integer']}, "
        f"min={info['value_min']}, max={info['value_max']})"
    )
    return level, ev


def _infer_feature_id_type(var_index: list[str]) -> tuple[str | None, str]:
    from src.workflows.manifest_data_validation import check_feature_id_type

    # declared arg only affects wording; detected_pattern is computed from the data.
    res = check_feature_id_type(var_index, "gene_symbol")
    pattern = res.get("detected_pattern", "probe_id_or_unknown")
    return _FEATURE_MAP.get(pattern), f"var.index looks like '{pattern}' — {res['message']}"


def _candidate_group_columns(obs) -> list[str]:
    """obs columns that look like usable grouping factors (2–12 categories)."""
    n = len(obs)
    out: list[str] = []
    for col in obs.columns:
        try:
            nu = int(obs[col].astype(str).nunique())
        except Exception:  # noqa: BLE001
            continue
        if 2 <= nu <= 12 and nu < n:
            out.append(str(col))
    return out


def _base_manifest(
    dataset_id: str,
    *,
    organism: str,
    modality: str,
    data_level: str | None,
    feature_id_type: str | None,
    group_columns: list[str],
    embedded_obs: bool,
    title: str | None,
) -> dict:
    return {
        "dataset_id": dataset_id,
        "title": title or f"{_TODO}: descriptive title for {dataset_id}",
        "accession": _TODO,
        "organism": organism,
        "modality": modality,
        "platform": _TODO,
        "data_level": data_level or _TODO,
        "feature_id_type": feature_id_type or _TODO,
        "expression_source": {"type": "local"},
        "metadata_source": {"type": "local", "embedded": embedded_obs},
        "group_columns": group_columns,
        "valid_workflows": [],
        "limitations": [
            "Manifest auto-drafted from the uploaded file — review every field before registering.",
        ],
    }


def draft_manifest(
    path: str | Path,
    *,
    dataset_id: str | None = None,
    title: str | None = None,
    organism: str = "human",
) -> DraftResult:
    """Inspect an upload and return a pre-filled manifest skeleton + a todo list.

    Parameters mirror the few prompts a UI would ask (``dataset_id``, ``title``,
    ``organism``); everything else is inferred from the file or left as ``TODO``.
    Raises ``ValueError`` for an unsupported file type or an unreadable file — the
    caller surfaces that to the uploader.
    """
    path = Path(path)
    kind = _kind_of(path.name)
    if kind is None:
        raise ValueError(
            f"Cannot draft a manifest for '{path.name}': supported types are "
            ".h5ad and flat matrices (.csv/.tsv/.txt, optionally .gz)."
        )
    dataset_id = dataset_id or _slug(path.name)

    inferred: dict[str, Any] = {}
    notes: list[str] = []

    if kind == "h5ad":
        import scanpy as sc

        adata = sc.read_h5ad(path)  # vetted loader — never exec
        X = adata.X
        if hasattr(X, "toarray"):
            X = X.toarray()
        X_flat = _cap_flat(X.flatten())
        var_index = [str(v) for v in adata.var.index]

        data_level, dl_ev = _infer_data_level(X_flat)
        feat, feat_ev = _infer_feature_id_type(var_index)
        group_cols = _candidate_group_columns(adata.obs)
        n_samples, n_features = int(adata.n_obs), int(adata.n_vars)

        modality = "sc_rnaseq" if n_samples > 2000 else "bulk_rnaseq"
        notes.append(
            f"modality guessed '{modality}' from {n_samples} obs — confirm "
            "(single-cell vs bulk changes the analysis path)."
        )
        inferred["group_columns"] = (
            group_cols,
            f"low-cardinality obs columns among {list(adata.obs.columns)}",
        )
        embedded_obs = True
    else:
        import pandas as pd

        name = path.name.lower()
        base = name[:-3] if name.endswith(".gz") else name
        sep = "\t" if base.endswith((".tsv", ".txt")) else ","
        df = pd.read_csv(path, index_col=0, sep=sep)  # vetted reader — never exec
        X_flat = _cap_flat(df.to_numpy(dtype=float).flatten())
        var_index = [str(c) for c in df.columns]

        data_level, dl_ev = _infer_data_level(X_flat)
        feat, feat_ev = _infer_feature_id_type(var_index)
        group_cols = []
        n_samples, n_features = int(df.shape[0]), int(df.shape[1])
        modality = "bulk_rnaseq"
        notes.append(
            "Flat matrix has no sample metadata — group_columns left empty. "
            "Re-upload as .h5ad with embedded obs to draft grouping/contrasts."
        )
        notes.append("Assumed orientation: samples as rows, genes as columns (index_col=0).")
        embedded_obs = False

    inferred["data_level"] = (data_level, dl_ev)
    inferred["feature_id_type"] = (feat, feat_ev)
    inferred["n_samples"] = (n_samples, "")
    inferred["n_features"] = (n_features, "")

    manifest = _base_manifest(
        dataset_id,
        organism=organism,
        modality=modality,
        data_level=data_level,
        feature_id_type=feat,
        group_columns=group_cols,
        embedded_obs=embedded_obs,
        title=title,
    )

    # Everything still needing a human: unfilled required fields + workflows.
    todo = [k for k, v in manifest.items() if v == _TODO]
    if not manifest["valid_workflows"]:
        todo.append("valid_workflows")
    if not manifest["group_columns"] and kind == "h5ad":
        notes.append("No obvious grouping column found — set group_columns manually.")

    return DraftResult(manifest=manifest, inferred=inferred, todo=todo, notes=notes)


def _slug(filename: str) -> str:
    stem = filename.lower()
    for suf in (".csv.gz", ".tsv.gz", ".txt.gz", ".h5ad", ".csv", ".tsv", ".txt"):
        if stem.endswith(suf):
            stem = stem[: -len(suf)]
            break
    keep = [c if (c.isalnum() or c == "_") else "_" for c in stem]
    slug = "".join(keep).strip("_") or "uploaded_dataset"
    return slug