"""``simurg methyl qc`` / ``simurg methyl convert-manifest`` — methylation control QC (spec MET-001 … MET-018).

Per sample (streamed, worker processes): control table, detection p-values, noob + dye-bias corrected β and M-values.
Per batch: ``methyl_qc_summary.tsv/.xlsx/.html``, ``pcs.tsv.gz``, ``probes.tsv``/``controls.tsv`` and the shared
``controls_raw.tsv``/``controls_qc.tsv`` (so ``simurg qc report`` accepts the folder).
"""

from __future__ import annotations

import datetime as _dt
import gzip
import hashlib
import html
import json
import math
import multiprocessing as mp
import os
import pickle
import tempfile
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np
import numpy.typing as npt

from simurg.core.atomic import check_outputs_absent
from simurg.core.errors import ExitCode, SimurgError
from simurg.core.log import Reporter
from simurg.core.provenance import RunManifest
from simurg.formats.idat import check_pair, read_idat
from simurg.formats.methyl_manifest import MethylManifest, read_methyl_manifest, write_probe_tables
from simurg.pipeline.qc import ControlSample, idat_metadata, write_control_tables, write_tsv
from simurg.qc.controls import ControlBead, canonical_id, compute_control_metrics, metrics_for
from simurg.qc.thresholds import Thresholds, parse_cli_pairs, resolve_thresholds
from simurg.signal import methyl as sig

SUMMARY = "methyl_qc_summary"
CONTROL_TYPES = ("STAINING", "EXTENSION", "HYBRIDIZATION", "TARGET REMOVAL", "NON-POLYMORPHIC", "SPECIFICITY I",
                 "SPECIFICITY II", "BISULFITE CONVERSION I", "BISULFITE CONVERSION II", "NEGATIVE", "NORM_A",
                 "NORM_C", "NORM_G", "NORM_T")  # fmt: skip
AVG_COLUMNS = ("avg_green_raw", "avg_red_raw", "avg_green_norm", "avg_red_norm")


@dataclass
class MethylQcOptions:
    out: Path
    manifest: Path
    idats: Path | None = None
    sheet: Path | None = None
    config: Path | None = None
    profile: str | None = "auto"
    thresholds: list[str] = field(default_factory=list)
    bg_offset: float | None = None
    bs_conv_red_bg_set: str = "U3-U5"
    pval_threshold: float = 0.05
    beta_offset: float = 0.0
    dye_bias: str = "linear"
    min_prop_passing: float | None = None
    pca_max_samples_full: int = 100
    pca_probes: int = 10_000
    seed: int = 42
    threads: int = 0
    force: bool = False
    on_duplicate_idat: str = "error"
    argv: list[str] = field(default_factory=list)


@dataclass
class _Job:
    sample_id: str
    grn: str | None
    red: str | None
    sheet_values: dict[str, str]
    out_dir: str


_W: dict[str, Any] = {}


def _init_from_file(path: str) -> None:
    with open(path, "rb") as fh:
        _W.update(pickle.load(fh))


def _write_gz_tsv(path: Path, header: str, lines: list[str]) -> None:
    tmp = path.with_suffix(path.suffix + ".part")
    with gzip.open(tmp, "wt", encoding="utf-8", newline="\n", compresslevel=6) as fh:
        fh.write(header + "\n")
        fh.write("\n".join(lines))
        if lines:
            fh.write("\n")
    tmp.replace(path)


def _g6(v: float) -> str:
    return "NA" if not math.isfinite(v) else f"{v:.6g}"


def _process(job: _Job) -> dict[str, Any]:
    man: MethylManifest = _W["man"]
    opts: MethylQcOptions = _W["opts"]
    beads: list[ControlBead] = _W["beads"]
    row: dict[str, Any] = {"sample_id": job.sample_id, "status": "OK", "fail_reason": ""}
    if not job.grn or not job.red:
        row.update(status="FAILED", fail_reason="W_METH_MISSING_SAMPLE: IDATs absent")
        return row
    step = "read"
    try:
        grn, red = read_idat(job.grn), read_idat(job.red)
        check_pair(grn, red)
        row["meta"] = idat_metadata(grn, red, job.sheet_values)
        row["sentrix_id"], row["sentrix_position"] = grn.barcode or red.barcode, grn.position or red.position
        row["scan_time"] = red.scan_time_iso or grn.scan_time_iso
        step = "signal assembly"
        raw = sig.assemble(grn, red, man)
        g, r = raw.ctrl_green, raw.ctrl_red
        qc = compute_control_metrics(beads, g, r, "methylation", _W["bg_offset"], opts.bs_conv_red_bg_set)
        row["control"] = ControlSample(job.sample_id, row["meta"], g, r, qc)
        step = "detection"
        pval = sig.detection_p(raw)
        step = "noob"
        corr, _, _ = sig.noob(raw)
        step = "dye bias"
        corr = sig.dye_bias_nonlinear(corr) if opts.dye_bias == "nonlinear" else sig.dye_bias_linear(corr)
        step = "beta"
        beta = sig.beta_values(corr, opts.beta_offset)
        mval = sig.m_values(beta)
        ana = _W["analytical"]
        row["n_probes"] = int(ana.sum())
        row["n_pass_detection"] = int((pval[ana] < opts.pval_threshold).sum())
        row["prop_probes_passing"] = row["n_pass_detection"] / row["n_probes"] if row["n_probes"] else math.nan
        row["avg_green_raw"], row["avg_red_raw"] = sig.average_signals(raw, ana)
        row["avg_green_norm"], row["avg_red_norm"] = sig.average_signals(corr, ana)
        row["pca"] = beta[_W["pca_idx"]].astype(np.float32)
        step = "write"
        d = Path(job.out_dir)
        d.mkdir(parents=True, exist_ok=True)
        _write_gz_tsv(
            d / f"{job.sample_id}_ctrl.tsv.gz",
            "address\ttype\tcolor_channel\tname\tprobe_id\tgreen\tred",
            [
                f"{c.address}\t{c.type}\t{c.color_channel}\t{c.name}\t{c.probe_id}\t{_g6(gv)}\t{_g6(rv)}"
                for c, gv, rv in zip(man.controls, g.tolist(), r.tolist(), strict=True)
            ],
        )
        _write_gz_tsv(
            d / f"{job.sample_id}_cgs.tsv.gz",
            "probe_id\tbeta\tm_value\tdetection_p",
            [
                f"{pid}\t{_g6(b)}\t{_g6(mv)}\t{_g6(pv)}"
                for pid, b, mv, pv in zip(man.probe_ids, beta.tolist(), mval.tolist(), pval.tolist(), strict=True)
            ],
        )
    except Exception as e:  # MET-010: this sample fails, the others continue
        code = e.code if isinstance(e, SimurgError) else "E_METH_PREP_FAILED"
        row.update(status="FAILED", fail_reason=f"{code}: Prep failed ({step}): {e}")
        row.pop("control", None)
    return row


def _pca(mat: npt.NDArray[np.float32]) -> npt.NDArray[np.float64] | None:
    """First two principal components (samples × 2) of a samples × probes matrix; None when not computable."""
    ok = np.isfinite(mat).all(axis=0)
    x = mat[:, ok].astype(np.float64)
    if x.shape[0] < 2 or x.shape[1] < 2:
        return None
    x -= x.mean(axis=0)
    gram = x @ x.T
    w, v = np.linalg.eigh(gram)
    order = np.argsort(w)[::-1][:2]
    pcs = v[:, order] * np.sqrt(np.maximum(w[order], 0))
    for j in range(pcs.shape[1]):  # deterministic sign: largest |loading| positive
        k = int(np.argmax(np.abs(pcs[:, j])))
        if pcs[k, j] < 0:
            pcs[:, j] = -pcs[:, j]
    if pcs.shape[1] < 2:
        pcs = np.column_stack([pcs, np.zeros(len(pcs))])
    return pcs


def _thresholds(opts: MethylQcOptions, man: MethylManifest) -> Thresholds:
    cli: dict[str, Any] = parse_cli_pairs(opts.thresholds)
    if opts.bg_offset is not None:
        cli["bgOffset"] = opts.bg_offset
    return resolve_thresholds(
        config=opts.config,
        profile=opts.profile,
        cli=cli,
        array_type="methylation",
        manifest_name=man.heading.get("Descriptor File Name", Path(man.path).name),
    )


def run_methyl_qc(opts: MethylQcOptions, rep: Reporter) -> int:
    from simurg.pipeline.inputs import check_paths, resolve_sample_set, sheet_values

    rm = RunManifest(command="methyl qc", argv=opts.argv)
    rep.run_id = rm.run_id
    check_paths([(opts.manifest, "methylation manifest"), (opts.idats, "IDAT folder"), (opts.sheet, "sheet")])
    try:
        sset = resolve_sample_set(opts.idats, opts.sheet, opts.on_duplicate_idat, rep, rm)
    except SimurgError as e:
        if e.code == "E_NO_INPUT":
            raise SimurgError("E_METH_NO_IDATS", f"No IDATs found ({e.message})") from None
        raise
    outputs = [opts.out / f"{SUMMARY}.{x}" for x in ("tsv", "xlsx", "html")] + [opts.out / "pcs.tsv.gz"]
    check_outputs_absent(outputs, force=opts.force)
    rep.info(f"reading methylation manifest {opts.manifest}")
    man = read_methyl_manifest(opts.manifest)
    opts.out.mkdir(parents=True, exist_ok=True)
    for p in write_probe_tables(man, opts.out):
        rm.add_output(p, "simurg.methyl_manifest/1")
    th = _thresholds(opts, man)
    for rec in th.as_records():
        rm.set_option(f"qc.{rec['key']}", rec["value"], rec["source"])
    beads = [ControlBead(c.type, c.color_channel, c.name, c.address, canonical_id(c.name)) for c in man.controls]
    types = {c.type.upper() for c in man.controls}
    if miss := [t for t in CONTROL_TYPES if t not in types]:
        rep.warn("W_METH_MISSING_CONTROL_TYPES", f"manifest lacks control type(s): {', '.join(miss)}")
    analytical = man.analytical()
    n = len(sset.samples)
    cand = np.flatnonzero(analytical)
    if n > opts.pca_max_samples_full and len(cand) > opts.pca_probes:
        cand = np.sort(np.random.default_rng(opts.seed).choice(cand, opts.pca_probes, replace=False))
    jobs = []
    for s in sset.samples:
        grn = str(s.pair.grn) if s.pair is not None and s.pair.grn is not None else None
        red = str(s.pair.red) if s.pair is not None and s.pair.red is not None else None
        if not (grn and red):
            rep.warn("W_METH_MISSING_SAMPLE", f"Missing sample: {s.sample_id}", sample_id=s.sample_id)
        jobs.append(_Job(s.sample_id, grn, red, sheet_values(s, sset.sheet), str(opts.out / "samples" / s.sample_id)))
    state = {"man": man, "opts": opts, "beads": beads, "bg_offset": th.bg_offset, "analytical": analytical,
             "pca_idx": cand}  # fmt: skip
    workers = max(1, min(opts.threads or (os.cpu_count() or 1), len(jobs) or 1))
    rep.info(f"methylation QC of {len(jobs)} sample(s) with {workers} worker(s)")
    if workers == 1:
        _W.update(state)
        try:
            rows = [_process(j) for j in jobs]
        finally:
            _W.clear()
    else:
        fd, path = tempfile.mkstemp(prefix=".simurg-workers-", suffix=".pkl", dir=opts.out)
        try:
            with os.fdopen(fd, "wb") as fh:
                pickle.dump(state, fh, protocol=5)
            ctx = mp.get_context("spawn")
            with ProcessPoolExecutor(workers, mp_context=ctx, initializer=_init_from_file, initargs=(path,)) as ex:
                rows = list(ex.map(_process, jobs))
        finally:
            Path(path).unlink(missing_ok=True)

    metrics = metrics_for("methylation")
    col_key = {m.column: m.key for m in metrics}
    missing_names: dict[str, list[str]] = {}
    for r in rows:
        if r["status"] != "OK":
            rep.error(
                r["fail_reason"].split(":", 1)[0], f"{r['sample_id']}: {r['fail_reason']}", sample_id=r["sample_id"]
            )
            continue
        q = r["control"].qc
        missing_names.update(q.missing)
        if not r["scan_time"]:
            rep.warn("W_IDAT_NO_SCAN_TIME", f"{r['sample_id']}: No scan time available", sample_id=r["sample_id"])
        fails = [m.column for m in metrics if m.column not in q.missing and th.get(m.key).fails(q.values[m.column])]
        if opts.min_prop_passing is not None and not r["prop_probes_passing"] >= opts.min_prop_passing:
            fails.append("prop_probes_passing")
        r["fail_codes"] = ";".join(fails)
        r["na_metrics"] = ";".join(q.na_metrics)
        r["pass_qc"] = 0 if fails else 1
        r.update({m.column: q.values[m.column] for m in metrics})
    if missing_names:
        ids = sorted({c for v in missing_names.values() for c in v})
        rep.warn("W_METH_MISSING_CONTROL_NAMES", f"manifest lacks control(s) {', '.join(ids)}; "
                 f"not evaluable: {', '.join(sorted(missing_names))}")  # fmt: skip

    meta_cols: list[str] = []
    for j in jobs:
        for k in j.sheet_values:
            if k not in meta_cols:
                meta_cols.append(k)
    cols = ["sentrix_id", "sentrix_position", "sample_id", *meta_cols, *(m.column for m in metrics), *AVG_COLUMNS,
            "scan_time", "n_probes", "n_pass_detection", "prop_probes_passing", "pass_qc", "fail_codes",
            "na_metrics", "status", "fail_reason"]  # fmt: skip
    out_rows = []
    for j, r in zip(jobs, rows, strict=True):
        out_rows.append({**j.sheet_values, **{c: r.get(c, "") for c in cols if c not in j.sheet_values}})
    summary_tsv = opts.out / f"{SUMMARY}.tsv"
    write_tsv(summary_tsv, cols, out_rows)
    rm.add_output(summary_tsv, "simurg.methyl_qc_summary/1")
    _write_xlsx(opts.out / f"{SUMMARY}.xlsx", cols, out_rows, th, col_key)

    ok_rows = [r for r in rows if r["status"] == "OK"]
    samples = [r["control"] for r in ok_rows]
    if samples:
        for p in write_control_tables(opts.out, samples, beads, th.bg_offset, rep, array_type="methylation"):
            rm.add_output(p, "simurg.controls/1")
    pcs = _pca(np.vstack([r["pca"] for r in ok_rows])) if len(ok_rows) >= 2 else None
    group_col = next((k for k in meta_cols if k.casefold() == "sample_group"), None)
    groups = {j.sample_id: (j.sheet_values.get(group_col, "") if group_col else "") for j in jobs}
    pc_rows = []
    if pcs is None:
        rep.warn("W_METH_NO_FEATURES", "No features, skipping PCA plot")
    else:
        for r, (p1, p2) in zip(ok_rows, pcs.tolist(), strict=True):
            pc_rows.append({"sample_id": r["sample_id"], "pc1": p1, "pc2": p2,
                            "sample_group": groups[r["sample_id"]] or "NA"})  # fmt: skip
    with gzip.open(opts.out / "pcs.tsv.gz", "wt", encoding="utf-8", newline="\n") as fh:
        fh.write("sample_id\tpc1\tpc2\tsample_group\n")
        for pr in pc_rows:
            fh.write(f"{pr['sample_id']}\t{pr['pc1']:.6g}\t{pr['pc2']:.6g}\t{pr['sample_group']}\n")
    _write_html(opts.out / f"{SUMMARY}.html", out_rows, pc_rows, th, metrics, man, len(cand))
    n_ok = len(ok_rows)
    rm.summary = {"samples": len(rows), "prepared": n_ok, "pass_qc": sum(r.get("pass_qc") == 1 for r in ok_rows)}
    rm.status = "completed"
    rm.write(opts.out / "run_manifest.json")
    rep.info(f"done: {n_ok}/{len(rows)} samples prepared, {rm.summary['pass_qc']} pass QC; summary {summary_tsv}")
    if n_ok == len(rows):
        return ExitCode.OK
    return ExitCode.PARTIAL if n_ok else ExitCode.ALL_FAILED


def _write_xlsx(
    path: Path, cols: list[str], rows: list[dict[str, Any]], th: Thresholds, col_key: dict[str, str]
) -> None:
    from openpyxl import Workbook
    from openpyxl.styles import Font, PatternFill

    bad = PatternFill("solid", fgColor="F5C28E")
    wb = Workbook()
    ws = wb.active
    ws.title = "Samples"
    ws.append(cols)
    for c in ws[1]:
        c.font = Font(bold=True)
    for r in rows:
        ws.append([_xl(r.get(c)) for c in cols])
        cur = ws.max_row
        for j, c in enumerate(cols, start=1):
            v = r.get(c)
            na = c in str(r.get("na_metrics", "")).split(";")
            metric_fails = c in col_key and isinstance(v, float) and not na and th.get(col_key[c]).fails(v)
            if metric_fails or (c == "pass_qc" and v == 0):
                ws.cell(cur, j).fill = bad
    ws.freeze_panes = "D2"
    ws.auto_filter.ref = ws.dimensions
    t2 = wb.create_sheet("Thresholds")
    t2.append(["key", "threshold", "direction", "source"])
    for k in col_key.values():
        t = th.get(k)
        t2.append([k, t.value, t.direction, t.source])
    tmp = path.with_suffix(".xlsx.part")
    wb.save(tmp)
    tmp.replace(path)


def _xl(v: Any) -> Any:
    if isinstance(v, float) and not math.isfinite(v):
        return "NaN" if math.isnan(v) else ("Inf" if v > 0 else "-Inf")
    return "" if v is None else v


def _write_html(
    path: Path,
    rows: list[dict[str, Any]],
    pcs: list[dict[str, Any]],
    th: Thresholds,
    metrics: list[Any],
    man: MethylManifest,
    n_pca_probes: int,
) -> None:
    """MET-013: self-contained page — summary table, detection histogram, PCA of β."""
    from simurg.pipeline.qc_report import _plotly

    ok = [r for r in rows if r.get("status") == "OK"]
    fin = [r for r in ok if math.isfinite(r["prop_probes_passing"])]
    hist = {
        "pass": [r["prop_probes_passing"] for r in fin if r.get("pass_qc") == 1],
        "fail": [r["prop_probes_passing"] for r in fin if r.get("pass_qc") != 1],
    }
    data = {
        "hist": hist,
        "pcs": pcs,
        "n_pca_probes": n_pca_probes,
    }
    head = ["sample_id", "pass_qc", "prop_probes_passing", "fail_codes", "na_metrics", "scan_time"]
    trs = "".join("<tr>" + "".join(f"<td>{html.escape(_cell(r.get(c)))}</td>" for c in head) + "</tr>" for r in rows)
    th_rows = "".join(
        f"<tr><td>{m.column}</td><td>{m.key}</td><td>{_cell(th.get(m.key).value)}</td></tr>" for m in metrics
    )
    from simurg import __version__
    from simurg.core.brand import BAR_CSS, CHART, FONT_STACK, favicon_link, report_bar

    generated = _dt.datetime.now(_dt.UTC).strftime("%Y-%m-%d %H:%M UTC")
    checksum = hashlib.sha256(json.dumps(rows, sort_keys=True, default=str).encode("utf-8")).hexdigest()
    bar = report_bar("Methylation control QC", __version__, "samples sha256", checksum)
    page = f"""<!doctype html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1"><title>Methylation QC · SimurgArray</title>
{favicon_link()}
<style>body{{font:14px {FONT_STACK};margin:0;color:#222;background:#fff}}.sa-main{{padding:16px}}
table{{border-collapse:collapse}}
td,th{{border:1px solid #ccc;padding:3px 6px;text-align:left}}.plots{{display:flex;flex-wrap:wrap;gap:16px}}
.plot{{width:min(560px,100%);height:380px}}.note{{color:#666}}{BAR_CSS}</style>
<script>{_plotly()}</script></head><body>
{bar}<div class="sa-main">
<p class="note">{len(rows)} samples · manifest {html.escape(Path(man.path).name)}
 · profile {html.escape(th.profile or "none")} · generated {generated}</p>
<div class="plots"><div id="hist" class="plot"></div><div id="pca" class="plot"></div></div>
<h2>Samples</h2><table><tr>{"".join(f"<th>{c}</th>" for c in head)}</tr>{trs}</table>
<h2>Thresholds</h2><table><tr><th>metric</th><th>key</th><th>threshold</th></tr>{th_rows}</table>
<script>const D={json.dumps(data, allow_nan=False, default=_json_num)};
Plotly.newPlot("hist",[{{x:D.hist.pass,type:"histogram",name:"pass",marker:{{color:"{CHART["BB"]}"}}}},
{{x:D.hist.fail,type:"histogram",name:"fail",marker:{{color:"{CHART["AA"]}"}}}}],
{{title:"Probes passing detection",barmode:"overlay",xaxis:{{title:"proportion"}},yaxis:{{title:"samples"}}}});
if(D.pcs.length){{const g=[...new Set(D.pcs.map(p=>p.sample_group))];
Plotly.newPlot("pca",g.map(k=>({{x:D.pcs.filter(p=>p.sample_group===k).map(p=>p.pc1),
y:D.pcs.filter(p=>p.sample_group===k).map(p=>p.pc2),text:D.pcs.filter(p=>p.sample_group===k).map(p=>p.sample_id),
mode:"markers",type:"scatter",name:k}})),{{title:"PCA of beta ("+D.n_pca_probes+" candidate probes)",
xaxis:{{title:"PC1"}},yaxis:{{title:"PC2"}}}});}}
else{{document.getElementById("pca").textContent="No features, PCA skipped";}}
</script></div></body></html>
"""
    tmp = path.with_suffix(".html.part")
    tmp.write_text(page, encoding="utf-8")
    tmp.replace(path)


def _cell(v: Any) -> str:
    if v is None:
        return ""
    if isinstance(v, float):
        return "NaN" if math.isnan(v) else f"{v:.4g}"
    return str(v)


def _json_num(v: Any) -> Any:
    if isinstance(v, np.floating | np.integer):
        return v.item()
    raise TypeError(type(v))


def run_convert_manifest(manifest: Path, out: Path, force: bool, rep: Reporter) -> int:
    """``simurg methyl convert-manifest`` (MET-003)."""
    check_outputs_absent([out / "probes.tsv", out / "controls.tsv"], force=force)
    man = read_methyl_manifest(manifest)
    probes, ctl = write_probe_tables(man, out)
    rep.info(f"{man.num_probes} probes -> {probes}; {len(man.controls)} controls -> {ctl}")
    return ExitCode.OK