#!/usr/bin/env python3
"""Phase P1 validation on a large data set (e.g. 1000 Genomes Omni2.5, 2141 samples with vendor GTCs).

For every GTC: parse, compare every field with Illumina BeadArrayFiles (BSD-2) and test the byte-identical
round-trip. For every IDAT pair: parse, pair-check and compare the address set with the manifest. BPM/EGT/CSV are
compared with BeadArrayFiles once. Writes a JSON summary; exit code 0 only if everything matched.

Usage: validate_p1_large.py --root DIR --bpm F --egt F [--csv F] --out summary.json [--workers N] [--limit N]
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import traceback
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from typing import Any

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "refs" / "BeadArrayFiles"))


def check_gtc(path: str) -> dict[str, Any]:
    from module.GenotypeCalls import GenotypeCalls  # type: ignore[import-not-found]

    from simurg.formats.gtc import encode_gtc, read_gtc

    res: dict[str, Any] = {"file": path, "ok": True, "diffs": []}
    try:
        g = read_gtc(path, check_write_complete=False)
        raw = Path(path).read_bytes()
        res["version"] = g.version
        res["autocall_version"] = g.autocall_version
        res["roundtrip_identical"] = encode_gtc(g) == raw
        if not res["roundtrip_identical"]:
            res["roundtrip_canonical_identical"] = encode_gtc(g, canonical=True) == raw
            res["size_ours_vs_file"] = [len(encode_gtc(g)), len(raw)]
        b = GenotypeCalls(path, check_write_complete=False)
        pairs = {
            "raw_x": (g.raw_x, b.get_raw_x_intensities()),
            "raw_y": (g.raw_y, b.get_raw_y_intensities()),
            "genotypes": (g.genotypes, np.array(b.get_genotypes(), dtype=np.uint8)),
            "scores": (g.genotype_scores, b.get_genotype_scores()),
        }
        if g.version >= 4:
            pairs["baf"] = (g.b_allele_freqs, b.get_ballele_freqs())
            pairs["lrr"] = (g.logr_ratios, b.get_logr_ratios())
        for k, (a, c) in pairs.items():
            if not np.array_equal(np.asarray(a), np.asarray(c), equal_nan=a.dtype.kind == "f"):
                res["diffs"].append(k)
        scalars = {
            "call_rate": (g.call_rate, b.get_call_rate()),
            "gender": (g.gender, b.get_gender()),
            "logr_dev": (g.logr_dev, b.get_logr_dev()),
            "gc10": (g.gc10, b.get_gc10()),
            "gc50": (g.gc50, b.get_gc50()),
            "num_calls": (g.num_calls, b.get_num_calls()),
            "sample_name": (g.sample_name, b.get_sample_name()),
            "base_calls": (g.base_call_strings(), [x.decode() for x in b.get_base_calls()]),
        }
        for k, (a, c) in scalars.items():
            same = (a == c) or (isinstance(a, float) and np.isnan(a) and np.isnan(c))
            if not same:
                res["diffs"].append(k)
        xf_b = [
            (t.offset_x, t.offset_y, t.scale_x, t.scale_y, t.shear, t.theta) for t in b.get_normalization_transforms()
        ]
        xf_g = [
            tuple(float(x[k]) for k in ("offset_x", "offset_y", "scale_x", "scale_y", "shear", "theta"))
            for x in g.normalization_transforms
        ]
        if [tuple(float(v) for v in t) for t in xf_b] != xf_g:
            res["diffs"].append("normalization_transforms")
        res["ok"] = not res["diffs"]
    except Exception as e:
        res.update(ok=False, error=f"{type(e).__name__}: {e}", trace=traceback.format_exc(limit=3))
    return res


def check_idat_pair(args: tuple[str, str]) -> dict[str, Any]:
    from simurg.formats.idat import check_pair, read_idat

    grn, red = args
    res: dict[str, Any] = {"grn": grn, "ok": True}
    try:
        g, r = read_idat(grn), read_idat(red)
        check_pair(g, r)
        res.update(
            n=g.n_addresses,
            chip=g.chip_type,
            midblock_itemsize=g.mid_block_itemsize,
            scan=g.scan_time_iso,
            name_ok=(g.barcode in Path(grn).name),
        )
    except Exception as e:
        res.update(ok=False, error=f"{type(e).__name__}: {e}")
    return res


_P2: dict[str, Any] = {}


def _p2_init(manifest: Any) -> None:
    _P2.clear()
    _P2["m"] = manifest


def check_raw_xy(args: tuple[str, str, str]) -> dict[str, Any]:
    """P2: raw X/Y + controls assembled from IDATs vs the arrays stored in the GTC (TOC 1000/1001/500/501)."""
    from simurg.formats.gtc import read_gtc
    from simurg.formats.idat import read_idat
    from simurg.signal.raw_xy import LocusMapCache, assemble_raw_xy

    gtc, grn, red = args
    res: dict[str, Any] = {"gtc": gtc, "ok": True}
    try:
        m = _P2["m"]
        cache = _P2.setdefault("cache", LocusMapCache(m))
        ref = read_gtc(gtc, check_write_complete=False)
        raw = assemble_raw_xy(read_idat(grn), read_idat(red), m, cache=cache)
        dx = raw.raw_x != ref.raw_x
        dy = raw.raw_y != ref.raw_y
        diff = dx | dy
        res.update(
            n_diff=int(diff.sum()),
            n_missing=raw.n_missing,
            vendor_zero=int(((ref.raw_x == 0) & (ref.raw_y == 0)).sum()),
            # plan §20 U1: would swapping channels explain the differences?
            swapped=int(((raw.raw_x == ref.raw_y) & (raw.raw_y == ref.raw_x) & diff).sum()),
            ours_missing_vendor_nonzero=int((raw.missing & ((ref.raw_x != 0) | (ref.raw_y != 0))).sum()),
            vendor_zero_ours_present=int((~raw.missing & (ref.raw_x == 0) & (ref.raw_y == 0)).sum()),
            controls_equal=bool(
                np.array_equal(raw.controls_x, ref.controls_x) and np.array_equal(raw.controls_y, ref.controls_y)
            ),
        )
        if res["n_diff"]:
            idx = np.flatnonzero(diff)[:5].tolist()
            res["examples"] = [
                {
                    "locus": m.names[i],
                    "assay": int(m.assay_types[i]),
                    "ours": [int(raw.raw_x[i]), int(raw.raw_y[i])],
                    "vendor": [int(ref.raw_x[i]), int(ref.raw_y[i])],
                }
                for i in idx
            ]
        res["ok"] = res["n_diff"] == 0 and res["controls_equal"]
    except Exception as e:
        res.update(ok=False, error=f"{type(e).__name__}: {e}")
    return res


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", required=True, type=Path)
    ap.add_argument("--bpm", required=True, type=Path)
    ap.add_argument("--egt", required=True, type=Path)
    ap.add_argument("--csv", type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--workers", type=int, default=os.cpu_count() or 1)
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--no-p2", action="store_true", help="skip the phase-P2 raw X/Y check")
    a = ap.parse_args()

    from simurg.formats.bpm import read_bpm
    from simurg.formats.discovery import walk_files
    from simurg.formats.egt import read_egt

    t0 = time.time()
    summary: dict[str, Any] = {"root": str(a.root), "started": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}

    gtcs = [str(p) for p in walk_files(a.root, (".gtc",))]
    idats = walk_files(a.root, (".idat", ".idat.gz"))
    if a.limit:
        gtcs = gtcs[: a.limit]
    grn = {p.name.lower().replace("_grn.idat", ""): p for p in idats if "_grn.idat" in p.name.lower()}
    red = {p.name.lower().replace("_red.idat", ""): p for p in idats if "_red.idat" in p.name.lower()}
    pairs = [(str(grn[k]), str(red[k])) for k in sorted(set(grn) & set(red))]
    if a.limit:
        pairs = pairs[: a.limit]
    summary["n_gtc"] = len(gtcs)
    summary["n_idat_pairs"] = len(pairs)
    summary["unpaired_idats"] = len(set(grn) ^ set(red))

    with ProcessPoolExecutor(a.workers) as ex:
        gres = list(ex.map(check_gtc, gtcs, chunksize=4))
        ires = list(ex.map(check_idat_pair, pairs, chunksize=8))
    summary["gtc_ok"] = sum(r["ok"] for r in gres)
    summary["gtc_roundtrip_identical"] = sum(bool(r.get("roundtrip_identical")) for r in gres)
    summary["gtc_versions"] = sorted({r.get("version") for r in gres if "version" in r})
    summary["gtc_autocall_versions"] = sorted({r.get("autocall_version") for r in gres if "autocall_version" in r})
    summary["gtc_failures"] = [r for r in gres if not r["ok"]][:50]
    summary["gtc_roundtrip_failures"] = [r for r in gres if r.get("roundtrip_identical") is False][:20]
    summary["idat_ok"] = sum(r["ok"] for r in ires)
    summary["idat_failures"] = [r for r in ires if not r["ok"]][:50]
    summary["idat_midblock_itemsizes"] = sorted({r.get("midblock_itemsize") for r in ires if r["ok"]})
    summary["idat_chip_types"] = sorted({r.get("chip") for r in ires if r["ok"]})

    # manifest / cluster file vs BeadArrayFiles
    import module  # type: ignore[import-not-found]

    m = read_bpm(a.bpm)
    b = module.BeadPoolManifest(str(a.bpm))
    summary["bpm"] = {
        "num_loci": m.num_loci,
        "names": m.names == b.names,
        "snps": m.snps == b.snps,
        "chroms": m.chroms == b.chroms,
        "map_infos": m.map_infos.tolist() == b.map_infos,
        "address_a": m.address_a.tolist() == b.addresses,
        "assay_types": m.assay_types.tolist() == b.assay_types,
        "normalization_lookups": m.normalization_lookups().tolist() == b.normalization_lookups,
        "n_bins": len(np.unique(m.norm_bin_ids())),
    }
    if not a.no_p2:
        import multiprocessing as mp
        import re

        key_re = re.compile(r"(\d+_r\d{2}c\d{2})", re.IGNORECASE)
        triples = []
        for g in gtcs:
            mk = key_re.search(Path(g).name)
            k = mk.group(1).lower() if mk else Path(g).stem.lower()
            if k in grn and k in red:
                triples.append((g, str(grn[k]), str(red[k])))
        summary["p2_matched_gtc_idat"] = len(triples)
        with ProcessPoolExecutor(
            a.workers, mp_context=mp.get_context("spawn"), initializer=_p2_init, initargs=(m.lite(),)
        ) as ex:
            p2 = list(ex.map(check_raw_xy, triples, chunksize=4))
        summary["p2_ok"] = sum(r["ok"] for r in p2)
        summary["p2_total_diff_loci"] = sum(r.get("n_diff", 0) for r in p2)
        summary["p2_swapped_loci"] = sum(r.get("swapped", 0) for r in p2)
        summary["p2_ours_missing_vendor_nonzero"] = sum(r.get("ours_missing_vendor_nonzero", 0) for r in p2)
        summary["p2_vendor_zero_ours_present"] = sum(r.get("vendor_zero_ours_present", 0) for r in p2)
        summary["p2_failures"] = [r for r in p2 if not r["ok"]][:30]
    e = read_egt(a.egt)
    with open(a.egt, "rb") as fh:
        be = module.ClusterFile.read_cluster_file(fh)
    bad = 0

    def same(a: float, b: float) -> bool:
        return bool(np.float32(a) == np.float32(b) or (np.isnan(a) and np.isnan(b)))

    for i, name in enumerate(e.names):
        rec = be.get_record(name)
        st = (rec.aa_cluster_stats, rec.ab_cluster_stats, rec.bb_cluster_stats)
        if not all(
            same(e.theta_mean[i, k], s.theta_mean)
            and same(e.r_mean[i, k], s.r_mean)
            and same(e.theta_dev[i, k], s.theta_dev)
            and same(e.r_dev[i, k], s.r_dev)
            and e.n[i, k] == s.N
            for k, s in enumerate(st)
        ):
            bad += 1
    summary["egt"] = {
        "num_records": e.num_records,
        "dbv": e.data_block_version,
        "extra_floats": e.extra_per_record,
        "record_mismatches": bad,
    }
    if a.csv:
        from simurg.formats.csv_manifest import compare_manifests, read_csv_manifest

        summary["bpm_csv_problems"] = compare_manifests(m, read_csv_manifest(a.csv, load_sequences=False))
    if gtcs:
        g0 = __import__("simurg.formats.gtc", fromlist=["read_gtc"]).read_gtc(gtcs[0], check_write_complete=False)
        summary["gtc_vs_bpm_num_snps"] = [g0.num_snps, m.num_loci]
        summary["gtc_transforms_vs_bins"] = [len(g0.normalization_transforms), summary["bpm"]["n_bins"]]

    summary["seconds"] = round(time.time() - t0, 1)
    p2_passed = a.no_p2 or (
        summary.get("p2_ok") == summary.get("p2_matched_gtc_idat", -1) and summary.get("p2_matched_gtc_idat", 0) > 0
    )
    passed = p2_passed and (
        summary["gtc_ok"] == len(gtcs)
        and summary["idat_ok"] == len(pairs)
        and all(v for k, v in summary["bpm"].items() if isinstance(v, bool))
        and summary["egt"]["record_mismatches"] == 0
        and not summary.get("bpm_csv_problems")
    )
    summary["passed"] = passed
    a.out.parent.mkdir(parents=True, exist_ok=True)
    a.out.write_text(json.dumps(summary, indent=2, default=str) + "\n")
    print(
        json.dumps(
            {
                k: summary.get(k)
                for k in (
                    "n_gtc",
                    "gtc_ok",
                    "gtc_roundtrip_identical",
                    "n_idat_pairs",
                    "idat_ok",
                    "p2_matched_gtc_idat",
                    "p2_ok",
                    "p2_total_diff_loci",
                    "passed",
                    "seconds",
                )
            }
        )
    )
    return 0 if passed else 1


if __name__ == "__main__":
    sys.exit(main())