SimurgArray/src/simurg/pipeline/normalize.py

"""``simurg normalize``: ingest output → per-sample normalization transforms and normalized X/Y, θ, R (phase P3).

Reads ``loci.parquet`` and every ``*.intensity.parquet`` under ``--in`` and writes, next to each input::

    <stem>.transforms.parquet   one row per normalization bin (GTC XForm fields)
    <stem>.normalized.parquet   norm_x, norm_y, theta, r per locus

plus ``normalize_summary.tsv`` and ``normalize_run_manifest.json`` in ``--in``.
"""

from __future__ import annotations

import multiprocessing as mp
import os
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np

from simurg import __version__
from simurg.core.atomic import atomic_write_text, check_outputs_absent
from simurg.core.errors import ExitCode, SimurgError
from simurg.core.log import Reporter
from simurg.core.provenance import RunManifest
from simurg.core.schemas import (
    INTENSITY_SCHEMA_ID,
    LOCI_SCHEMA_ID,
    normalized_table,
    read_metadata,
    read_table,
    transforms_table,
    write_parquet,
)
from simurg.formats.discovery import walk_files
from simurg.signal.normalization.apply import normalize, polar
from simurg.signal.normalization.fit import VERSIONS, fit_transforms

SUFFIX = ".intensity.parquet"


@dataclass
class NormalizeOptions:
    indir: Path
    version: str = "1.2.0"
    threads: int = 0
    force: bool = False
    argv: list[str] = field(default_factory=list)


_STATE: dict[str, Any] = {}


def _init(bins: np.ndarray, lookup: np.ndarray, version: str) -> None:
    _STATE.update(bins=bins, lookup=lookup, version=version)


def _outputs(p: Path) -> tuple[Path, Path]:
    stem = p.name[: -len(SUFFIX)]
    return p.with_name(f"{stem}.transforms.parquet"), p.with_name(f"{stem}.normalized.parquet")


def _process(path: str) -> dict[str, Any]:
    p = Path(path)
    row: dict[str, Any] = {"file": path, "status": "OK", "fail_reason": ""}
    try:
        t = read_table(p, INTENSITY_SCHEMA_ID)
        md = read_metadata(p)
        row["sample_id"] = md.get("sample_id", "")
        raw_x = t.column("raw_x").to_numpy().astype(np.uint16)
        raw_y = t.column("raw_y").to_numpy().astype(np.uint16)
        bins, lookup, version = _STATE["bins"], _STATE["lookup"], _STATE["version"]
        if len(raw_x) != len(bins):
            raise SimurgError("E_BPM_FORMAT", f"{p}: {len(raw_x)} loci but the manifest has {len(bins)}")
        xforms = fit_transforms(raw_x, raw_y, bins, version)
        nx, ny = normalize(raw_x, raw_y, lookup, xforms)
        theta, r = polar(nx, ny)
        uniq, counts = np.unique(bins, return_counts=True)
        rows = []
        for k, (bid, cnt, xf) in enumerate(zip(uniq, counts, xforms, strict=True)):
            rows.append(
                {
                    "bin_index": k,
                    "norm_bin_id": int(bid),
                    "n_loci": int(cnt),
                    "version": xf.version,
                    **{
                        f: float(getattr(xf, f))
                        for f in (
                            "offset_x",
                            "offset_y",
                            "scale_x",
                            "scale_y",
                            "shear",
                            "theta",
                            "cvx",
                            "cvy",
                            "nn12",
                            "rr12",
                            "taa",
                            "tbb",
                        )
                    },
                }
            )
        meta = {
            "sample_id": md.get("sample_id", ""),
            "sample_name": md.get("sample_name", ""),
            "normalization_version": version,
            "simurg_version": __version__,
            "source_sha256_grn": md.get("grn_sha256", ""),
            "source_sha256_red": md.get("red_sha256", ""),
        }
        tp, np_ = _outputs(p)
        write_parquet(transforms_table(rows, meta), tp)
        write_parquet(normalized_table(nx, ny, theta, r, meta), np_)
        row["n_bins"] = len(rows)
    except SimurgError as e:
        row.update(status="FAILED", fail_reason=f"{e.code}: {e.message}", error_code=e.code)
    except Exception as e:
        row.update(status="FAILED", fail_reason=f"E_INTERNAL: {type(e).__name__}: {e}", error_code="E_INTERNAL")
    return row


def run_normalize(opts: NormalizeOptions, rep: Reporter) -> int:
    if opts.version not in VERSIONS:
        raise SimurgError("E_CONFIG_KEY", f"unknown normalization version {opts.version!r}; use {VERSIONS}")
    loci = opts.indir / "loci.parquet"
    if not loci.exists():
        raise SimurgError("E_FILE_NOT_FOUND", f"{loci} not found — run `simurg ingest` first", path=str(loci))
    rm = RunManifest(command="normalize", argv=opts.argv)
    rep.run_id = rm.run_id
    rm.set_option("normalization_version", opts.version, "cli")
    lt = read_table(loci, LOCI_SCHEMA_ID)
    bins = lt.column("norm_bin_id").to_numpy().astype(np.uint8)
    lookup = lt.column("norm_lookup").to_numpy().astype(np.int32)
    files = [str(p) for p in walk_files(opts.indir, (SUFFIX,))]
    if not files:
        raise SimurgError("E_NO_INPUT", f"no *{SUFFIX} files under {opts.indir}")
    check_outputs_absent([q for f in files for q in _outputs(Path(f))], force=opts.force)

    workers = max(1, min(opts.threads or (os.cpu_count() or 1), len(files)))
    rep.info(f"normalizing {len(files)} sample(s) with version {opts.version}, {workers} worker(s)")
    if workers == 1:
        _init(bins, lookup, opts.version)
        results = [_process(f) for f in files]
    else:
        with ProcessPoolExecutor(
            workers, mp_context=mp.get_context("spawn"), initializer=_init, initargs=(bins, lookup, opts.version)
        ) as ex:
            results = list(ex.map(_process, files))
    for r in results:
        if r["status"] != "OK":
            rep.error(r.get("error_code", "E_INTERNAL"), f"{r['file']}: {r['fail_reason']}")
            rm.failed_samples.append({"sample_id": r.get("sample_id", ""), "code": r.get("error_code", "")})
    cols = ["sample_id", "status", "fail_reason", "n_bins", "file"]
    lines = ["\t".join(cols)] + ["\t".join(str(r.get(c, "")) for c in cols) for r in results]
    atomic_write_text(opts.indir / "normalize_summary.tsv", "\n".join(lines) + "\n")
    ok = sum(r["status"] == "OK" for r in results)
    rm.summary = {"samples": len(results), "succeeded": ok, "warnings": rep.warning_counts()}
    rm.status = "completed"
    rm.write(opts.indir / "normalize_run_manifest.json")
    rep.info(f"done: {ok}/{len(results)} samples normalized")
    if ok == len(results):
        return ExitCode.OK
    return ExitCode.PARTIAL if ok else ExitCode.ALL_FAILED