SimurgArray/src/simurg/formats/discovery.py

"""IDAT / GTC discovery and sample-set resolution — spec IN-016, IN-020 … IN-023, IN-030, CLI-006.

Directory listing order never influences results: every path list is sorted bytewise (``os.fsencode``).
"""

from __future__ import annotations

import os
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path

from simurg.core.errors import SimurgError
from simurg.formats.idat import parse_idat_filename, parse_run_time, read_idat
from simurg.formats.samplesheet import SampleSheet


@dataclass
class IdatPair:
    barcode: str
    position: str
    grn: Path | None
    red: Path | None

    @property
    def sample_id(self) -> str:
        return f"{self.barcode}_{self.position}"

    @property
    def complete(self) -> bool:
        return self.grn is not None and self.red is not None

    @property
    def folder(self) -> Path:
        p = self.grn or self.red
        assert p is not None
        return p.parent


@dataclass
class Discovery:
    root: Path
    pairs: dict[str, IdatPair]  # sample_id -> pair (complete or not)
    warnings: list[tuple[str, str]] = field(default_factory=list)


def _bytewise(paths: list[Path]) -> list[Path]:
    return sorted(paths, key=lambda p: os.fsencode(str(p)))


def walk_files(root: Path, suffixes: tuple[str, ...]) -> list[Path]:
    """Recursive listing following symlinks with cycle detection; sorted bytewise."""
    found: list[Path] = []
    seen_dirs: set[tuple[int, int]] = set()
    stack = [root]
    while stack:
        d = stack.pop()
        try:
            st = d.stat()
        except OSError:
            continue
        key = (st.st_dev, st.st_ino)
        if key in seen_dirs:
            continue
        seen_dirs.add(key)
        try:
            entries = list(os.scandir(d))
        except OSError:
            continue
        for e in entries:
            p = Path(e.path)
            try:
                if e.is_dir(follow_symlinks=True):
                    stack.append(p)
                elif e.is_file(follow_symlinks=True) and e.name.lower().endswith(suffixes):
                    found.append(p)
            except OSError:
                continue
    return _bytewise(found)


def _scan_time_key(p: Path) -> str:
    try:
        return parse_run_time(read_idat(p, load_arrays=False).scan_time_raw)
    except SimurgError:
        return ""


def discover_idats(root: str | Path, *, on_duplicate: str = "error") -> Discovery:
    """Find ``<barcode>_<pos>_{Grn,Red}.idat[.gz]`` under ``root`` and pair them (IN-020 … IN-022)."""
    rootp = Path(root)
    if not rootp.is_dir():
        raise SimurgError("E_FILE_NOT_FOUND", f"IDAT folder not found: {rootp}", path=str(rootp))
    by_key: dict[tuple[str, str, str], list[Path]] = defaultdict(list)
    for p in walk_files(rootp, (".idat", ".idat.gz")):
        parsed = parse_idat_filename(p)
        if parsed is None:
            continue
        by_key[parsed].append(p)

    warnings: list[tuple[str, str]] = []
    # group channel files by folder so that "the same sample in two folders" is a duplicate pair
    per_sample: dict[tuple[str, str], dict[Path, dict[str, Path]]] = defaultdict(lambda: defaultdict(dict))
    for (barcode, pos, ch), paths in by_key.items():
        for p in paths:
            per_sample[(barcode, pos)][p.parent][ch] = p

    pairs: dict[str, IdatPair] = {}
    duplicates: dict[str, list[Path]] = {}
    for (barcode, pos), folders in sorted(per_sample.items()):
        candidates = [
            IdatPair(barcode, pos, chans.get("Grn"), chans.get("Red"))
            for _, chans in sorted(folders.items(), key=lambda kv: os.fsencode(str(kv[0])))
        ]
        sid = f"{barcode}_{pos}"
        if len(candidates) > 1:
            if on_duplicate == "error":
                duplicates[sid] = [p for c in candidates for p in (c.grn, c.red) if p is not None]
                continue
            if on_duplicate == "newest":
                # candidates are sorted bytewise; max() keeps the first (smallest path) among equal scan times
                chosen = max(candidates, key=lambda c: _scan_time_key(c.grn or c.red))  # type: ignore[arg-type]
            else:  # "first": bytewise-smallest path
                chosen = candidates[0]
            warnings.append(
                (
                    "W_DUP_IDAT_RESOLVED",
                    f"{sid}: {len(candidates)} IDAT pairs found; using {chosen.folder} ({on_duplicate})",
                )
            )
        else:
            chosen = candidates[0]
        if not chosen.complete:
            have = "Grn" if chosen.grn else "Red"
            warnings.append(("W_IDAT_UNPAIRED", f"{sid}: only the {have} IDAT was found in {chosen.folder}"))
        pairs[sid] = chosen
    if duplicates:
        lines = [f"{sid}: " + ", ".join(str(p) for p in ps) for sid, ps in sorted(duplicates.items())]
        raise SimurgError(
            "E_DUP_IDAT",
            f"{len(duplicates)} sample(s) have more than one IDAT pair:\n  " + "\n  ".join(lines),
            hint="remove the extra copies or pass --on-duplicate-idat first|newest",
        )
    return Discovery(rootp, pairs, warnings)


@dataclass
class ResolvedSample:
    sample_id: str
    sample_name: str  # provisional: sheet Sample_Name or sample_id; IDAT field 406 is applied on read (IN-010)
    pair: IdatPair | None
    sheet_row: object | None
    status: str  # "READY" or "MISSING_INPUT"


def resolve_input_folder(
    cli_folder: str | Path | None, sheet: SampleSheet | None
) -> tuple[Path, list[tuple[str, str]]]:
    """Apply the input-set rule (CLI-006) and RootFolder precedence (IN-008)."""
    notes: list[tuple[str, str]] = []
    if cli_folder is None and sheet is None:
        raise SimurgError("E_NO_INPUT", "no input: give an IDAT folder, a sample sheet, or both")
    if cli_folder is not None:
        folder = Path(cli_folder)
        if sheet is not None and sheet.root_folder is not None and sheet.root_folder.resolve() != folder.resolve():
            notes.append(
                ("W_SHEET_ROOTFOLDER_OVERRIDDEN", f"CLI folder {folder} overrides sheet RootFolder {sheet.root_folder}")
            )
        return folder, notes
    assert sheet is not None
    if sheet.root_folder is None:
        raise SimurgError("E_NO_INPUT_FOLDER", f"{sheet.path}: sample sheet has no RootFolder and no folder was given")
    return sheet.root_folder, notes


def resolve_samples(disc: Discovery, sheet: SampleSheet | None) -> tuple[list[ResolvedSample], list[tuple[str, str]]]:
    """Match discovered IDAT pairs with sheet rows (IN-010, IN-016). Returns samples in a deterministic order."""
    notes: list[tuple[str, str]] = []
    out: list[ResolvedSample] = []
    if sheet is None:
        for sid in sorted(disc.pairs):
            found = disc.pairs[sid]
            status = "READY" if found.complete else "MISSING_INPUT"
            out.append(ResolvedSample(sid, sid, found, None, status))
        return out, notes
    listed = set()
    for row in sheet.rows:
        listed.add(row.sample_id)
        pair = disc.pairs.get(row.sample_id)
        status = "READY" if pair is not None and pair.complete else "MISSING_INPUT"
        name = row.sample_name or row.sample_id
        out.append(ResolvedSample(row.sample_id, name, pair, row, status))
    extra = sorted(set(disc.pairs) - listed)
    if extra:
        notes.append(("I_UNLISTED_IDATS", f"{len(extra)} IDAT pair(s) not listed in the sample sheet are ignored"))
    return out, notes


def discover_gtcs(root: str | Path) -> dict[str, Path]:
    """``*.gtc`` under ``root`` keyed by file-name stem (IN-030); duplicate keys → ``E_DUP_GTC``."""
    root = Path(root)
    if not root.exists():
        raise SimurgError("E_FILE_NOT_FOUND", f"GTC folder not found: {root}", path=str(root))
    out: dict[str, Path] = {}
    for p in walk_files(root, (".gtc",)):
        key = p.name[: -len(".gtc")]
        if key in out:
            raise SimurgError("E_DUP_GTC", f"GTC files {out[key]} and {p} resolve to the same sample {key!r}")
        out[key] = p
    if not out:
        raise SimurgError("E_NO_INPUT", f"no GTC files found under {root}")
    return out