#!/usr/bin/env python3
"""Compare `simurg methyl qc` numerics with the sesame oracle (scripts/v2/sesame_oracle.R) on the same probes.

python scripts/v2/methyl_oracle_compare.py --manifest M.csv --sesame DIR --out summary.json PREFIX...
"""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

import numpy as np

from simurg.formats.idat import read_idat
from simurg.formats.methyl_manifest import read_methyl_manifest
from simurg.signal import methyl as sig


def _q(a: np.ndarray) -> dict[str, float]:
    a = a[np.isfinite(a)]
    if not len(a):
        return {}
    return {k: float(np.quantile(a, q)) for k, q in (("median", 0.5), ("p99", 0.99), ("max", 1.0))}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--sesame", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("prefixes", nargs="+")
    a = ap.parse_args()
    man = read_methyl_manifest(a.manifest)
    pos = {p: i for i, p in enumerate(man.probe_ids)}
    res = {}
    for prefix in a.prefixes:
        # min_beads=0: the oracle runs without bead-count masking (see sesame_oracle.R)
        s = sig.assemble(read_idat(prefix + "_Grn.idat"), read_idat(prefix + "_Red.idat"), man, min_beads=0)
        rows = list(csv.DictReader((a.sesame / (Path(prefix).name + ".tsv")).open(), delimiter="\t"))
        idx = np.array([pos[r["probe_id"]] for r in rows])

        def col(name: str, rows: list[dict[str, str]] = rows) -> np.ndarray:
            return np.array([float(r[name]) if r[name] not in ("NA", "") else np.nan for r in rows])

        mask_s = np.array([r["mask"] == "TRUE" for r in rows])
        variants = {}
        for neg in (True, False):
            p = sig.detection_p(s, combine_neg=neg)
            n, fg, fr = sig.noob(s, combine_neg=neg)
            variants[neg] = (p, n, fg, fr, sig.beta_values(n))
        keep = ~mask_s & ~s.masked[idx]
        sp, sm, su, sb = col("pval"), col("m_noob"), col("u_noob"), col("beta_noob")
        ana = man.analytical()[idx] & keep
        out: dict[str, object] = {
            "probes": len(rows),
            "masked_sesame": int(mask_s.sum()),
            "masked_simurg": int(s.masked.sum()),
        }
        for neg, (p, n, fg, fr, b) in variants.items():
            op, om, ou, ob = p[idx], n.m[idx], n.u[idx], b[idx]
            out[f"combine_neg={neg}"] = {
                "noob_fit": {"green": vars(fg), "red": vars(fr)},
                "pval_abs_diff": _q(np.abs(op - sp)[keep]),
                "pval_identical_fraction": float((op[keep] == sp[keep]).mean()),
                "detection_call_agreement_0.05": float(((op < 0.05) == (sp < 0.05))[keep].mean()),
                "prop_passing": {"simurg": float((op[ana] < 0.05).mean()), "sesame": float((sp[ana] < 0.05).mean())},
                "noob_m_rel_diff": _q(np.abs(om - sm)[keep] / np.maximum(np.abs(sm[keep]), 1)),
                "noob_u_rel_diff": _q(np.abs(ou - su)[keep] / np.maximum(np.abs(su[keep]), 1)),
                "beta_abs_diff": _q(np.abs(ob - sb)[keep]),
                "beta_corr": float(np.corrcoef(ob[keep & np.isfinite(sb)], sb[keep & np.isfinite(sb)])[0, 1]),
            }
        res[Path(prefix).name] = out
    a.out.write_text(json.dumps(res, indent=1))
    print(json.dumps(res, indent=1))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())