#!/usr/bin/env python3
"""Content digests of SimurgArray outputs, for comparing runs across machines, architectures and installs (P10).

GTC: every numeric field and sample statistic, NaNs canonicalised (their sign/payload is platform-specific, K-014);
free-text metadata (dates, software version) is left out. VCF(.gz): all data lines, header lines excluded.
TSV: the listed columns of sample_summary.tsv.

    python scripts/p10/digest.py RUN_DIR [--out digest.json]
"""

from __future__ import annotations

import argparse
import gzip
import hashlib
import json
import sys
from pathlib import Path

import numpy as np

from simurg.formats.gtc import read_gtc

GTC_FIELDS = ("normalization_transforms", "controls_x", "controls_y", "raw_x", "raw_y", "genotypes", "base_calls",
              "genotype_scores", "b_allele_freqs", "logr_ratios", "percentiles_x", "percentiles_y", "sample_data",
              "call_rate", "gender", "logr_dev", "gc10")  # fmt: skip
SUMMARY_COLUMNS = ("sample_id", "status", "autosomal_call_rate", "call_rate", "logr_dev", "sex_est", "gc10",
                   "gc50", "n_calls", "n_nocalls", "het_rate", "contamination", "functional_status", "control_status",
                   "flags")  # fmt: skip


def _canon(a: np.ndarray) -> bytes:
    if a.dtype.names:  # structured: canonicalise float members
        a = a.copy()
        for n in a.dtype.names:
            if np.issubdtype(a.dtype[n], np.floating):
                v = a[n]
                v[np.isnan(v)] = np.nan
        return a.tobytes()
    if np.issubdtype(a.dtype, np.floating):
        a = np.where(np.isnan(a), np.array(np.nan, a.dtype), a)
    return np.ascontiguousarray(a).tobytes()


def gtc_digest(path: Path) -> dict[str, str]:
    g = read_gtc(path)
    out = {}
    for f in GTC_FIELDS:
        v = getattr(g, f)
        b = _canon(np.asarray(v)) if not isinstance(v, bytes) else v
        out[f] = hashlib.sha256(b).hexdigest()[:16]
    return out


def vcf_digest(path: Path) -> str:
    h = hashlib.sha256()
    with gzip.open(path, "rt") if path.suffix == ".gz" else open(path) as fh:
        for line in fh:
            if not line.startswith("#"):
                h.update(line.encode())
    return h.hexdigest()[:16]


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("run", type=Path)
    ap.add_argument("--out", type=Path)
    a = ap.parse_args()
    d: dict[str, object] = {}
    d["gtc"] = {p.stem: gtc_digest(p) for p in sorted(a.run.rglob("*.gtc"))}
    d["vcf"] = {p.name: vcf_digest(p) for p in sorted(a.run.rglob("*.vcf.gz"))}
    summ = {}
    for p in sorted(a.run.rglob("sample_summary.tsv")):
        rows = [ln.rstrip("\n").split("\t") for ln in p.open()]
        idx = [rows[0].index(c) for c in SUMMARY_COLUMNS if c in rows[0]]
        summ[str(p.relative_to(a.run))] = hashlib.sha256(
            "\n".join("\t".join(r[i] for i in idx) for r in rows).encode()
        ).hexdigest()[:16]
    d["summary"] = summ
    text = json.dumps(d, indent=1, sort_keys=True)
    (a.out.write_text(text) if a.out else sys.stdout.write(text + "\n"))
    return 0


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