#!/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


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)
    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())),
    }
    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)
    passed = (
        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[k]
                for k in ("n_gtc", "gtc_ok", "gtc_roundtrip_identical", "n_idat_pairs", "idat_ok", "passed", "seconds")
            }
        )
    )
    return 0 if passed else 1


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