SimurgArray/src/simurg/formats/methyl_manifest.py

"""Methylation CSV manifest (EPIC v1/v2, MSA and custom designs) — spec MET-002, MET-003; plan §6.3.

Sections: ``[Heading]`` (key,value), ``[Assay]`` (header row + one row per probe; needs ``IlmnID``, ``AddressA_ID``,
``AddressB_ID``, ``Color_Channel``) and ``[Controls]`` (``Address,Type,Color_Channel,Name``, no header row; the
vendor files pad lines with empty fields). Infinium I probes have an ``AddressB_ID``; the colour channel is read
from ``Color_Channel`` (``Red``/``Grn``), Infinium II probes leave both empty.
[DOC: Illumina methylation manifest column descriptions; layout of the open-source toolkits' manifests (MET-003)]
"""

from __future__ import annotations

import csv
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np
import numpy.typing as npt

from simurg.core.errors import FormatError
from simurg.formats.csv_manifest import _open_text

REQUIRED = ("IlmnID", "AddressA_ID", "AddressB_ID", "Color_Channel")


@dataclass
class MethylControl:
    address: int
    type: str
    color_channel: str
    name: str

    @property
    def probe_id(self) -> str:
        return f"ctl_{self.address}"


@dataclass
class MethylManifest:
    path: str
    heading: dict[str, str]
    probe_ids: list[str]
    u: npt.NDArray[np.int64]  # AddressA_ID
    m: npt.NDArray[np.int64]  # AddressB_ID; 0 for Infinium II
    col: npt.NDArray[np.str_]  # "R" / "G" for Infinium I, "" for Infinium II
    controls: list[MethylControl] = field(default_factory=list)

    @property
    def num_probes(self) -> int:
        return len(self.probe_ids)

    @property
    def type_i(self) -> npt.NDArray[np.bool_]:
        return np.asarray(self.m != 0, dtype=np.bool_)

    def analytical(self) -> npt.NDArray[np.bool_]:
        """CG and SNP probes (MET-011): IDs starting with ``cg`` or ``rs``."""
        return np.array([p[:2] in ("cg", "rs") for p in self.probe_ids], dtype=np.bool_)


def _address(s: str) -> int:
    s = s.strip()
    return int(s) if s else 0


def _channel(s: str, *, infinium_i: bool) -> str:
    c = s.strip().casefold()
    if not infinium_i:
        return ""
    if c in ("red", "r"):
        return "R"
    if c in ("grn", "green", "g"):
        return "G"
    return "?"


def read_methyl_manifest(path: str | Path) -> MethylManifest:
    """Parse a methylation CSV manifest (plain or gzipped). Raises the MET-002 errors."""
    p = Path(path)
    spath = str(p)
    heading: dict[str, str] = {}
    header: list[str] | None = None
    idx: dict[str, int] = {}
    ids: list[str] = []
    u: list[int] = []
    m: list[int] = []
    col: list[str] = []
    controls: list[MethylControl] = []
    section = None
    saw_assay = saw_controls = False
    with _open_text(p) as fh:
        for row in csv.reader(fh):
            if not row or not any(c.strip() for c in row):
                continue
            first = row[0].strip()
            if first.startswith("[") and first.endswith("]"):
                section = first[1:-1].strip().casefold()
                saw_assay |= section == "assay"
                saw_controls |= section == "controls"
                continue
            if section == "heading":
                heading[first] = row[1].strip() if len(row) > 1 else ""
            elif section == "assay":
                if header is None:
                    header = [c.strip() for c in row]
                    idx = {c: i for i, c in enumerate(header)}
                    miss = [c for c in REQUIRED if c not in idx]
                    if miss:
                        msg = f"{spath}: Missing required columns: {', '.join(miss)}"
                        raise FormatError("E_METH_MISSING_COLUMNS", msg, path=spath)
                    iu, im, ic, ii = idx["AddressA_ID"], idx["AddressB_ID"], idx["Color_Channel"], idx["IlmnID"]
                    continue
                ids.append(row[ii].strip())
                a, b = _address(row[iu]), _address(row[im]) if im < len(row) else 0
                u.append(a)
                m.append(b)
                col.append(_channel(row[ic] if ic < len(row) else "", infinium_i=b != 0))
            elif section == "controls":
                fields = list(row)
                while fields and not fields[-1].strip():
                    fields.pop()
                if len(fields) < 4:  # vendor EPIC files add a 5th field ("AVG") to some NEGATIVE lines (K-026)
                    raise FormatError(
                        "E_METH_CONTROLS_FORMAT",
                        f"{spath}: Controls not formatted correctly. Must have 4 columns "
                        f"(Address,Type,Color_Channel,Name): {','.join(fields)[:80]!r}",
                        path=spath,
                    )
                addr, typ, colour, name = (f.strip() for f in fields[:4])
                controls.append(MethylControl(_address(addr), typ, colour, name))
    if not saw_assay or header is None:
        raise FormatError("E_METH_NO_ASSAY", f"{spath}: Probe section not found", path=spath)
    if not saw_controls or not controls:
        raise FormatError("E_METH_NO_CONTROLS", f"{spath}: Missing control probes in manifest", path=spath)
    return MethylManifest(
        path=spath,
        heading=heading,
        probe_ids=ids,
        u=np.asarray(u, dtype=np.int64),
        m=np.asarray(m, dtype=np.int64),
        col=np.asarray(col, dtype="<U1"),
        controls=controls,
    )


def write_probe_tables(man: MethylManifest, out: Path) -> tuple[Path, Path]:
    """MET-003: ``probes.tsv`` (probe_id, u, m, col; ``NA`` for Infinium II) and ``controls.tsv``."""
    out.mkdir(parents=True, exist_ok=True)
    probes, ctl = out / "probes.tsv", out / "controls.tsv"
    with probes.open("w", encoding="utf-8", newline="\n") as fh:
        fh.write("probe_id\tu\tm\tcol\n")
        for pid, a, b, c in zip(man.probe_ids, man.u.tolist(), man.m.tolist(), man.col.tolist(), strict=True):
            fh.write(f"{pid}\t{a}\t{b if b else 'NA'}\t{c or 'NA'}\n")
    with ctl.open("w", encoding="utf-8", newline="\n") as fh:
        fh.write("address\ttype\tcolor_channel\tname\tprobe_id\n")
        for c in man.controls:
            fh.write(f"{c.address}\t{c.type}\t{c.color_channel}\t{c.name}\t{c.probe_id}\n")
    return probes, ctl