#!/usr/bin/env python3
"""Phase P4 validation against vendor GTCs: genotype every IDAT pair and compare the produced GTC with the vendor's,
field by field and as whole files (after aligning free-text metadata that cannot match).

Usage: validate_p4_large.py --root DIR --bpm F --egt-dir DIR --out summary.json
                            [--norm-version 1.1.2] [--sex-version 1] [--workers N] [--limit N]
"""

from __future__ import annotations

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

import numpy as np

NEG_NAN = np.frombuffer(np.uint32(0xFFC00000).tobytes(), dtype=np.float32)[0]
ALIGN = (
    "cluster_file",
    "snp_manifest",
    "autocall_date",
    "autocall_version",
    "imaging_date",
    "sample_name",
    "sample_plate",
    "sample_well",
    "scanner_data",
)
_W: dict[str, Any] = {}


def _init(bpm: str, egt_dir: str, norm: str, sex_version: int, preset: str) -> None:
    _W.update(bpm=bpm, egt_dir=egt_dir, norm=norm, sex_version=sex_version, preset=preset, ctx={})


def _ctx(cluster_field: str) -> Any:
    from simurg.formats.bpm import read_bpm
    from simurg.formats.egt import read_egt
    from simurg.genotype.clusters import align_clusters
    from simurg.genotype.sample import manifest_context

    name = re.sub(r"\(MD5Checksum=.*\)$", "", cluster_field).strip()
    if name not in _W["ctx"]:
        if "manifest" not in _W:
            _W["manifest"] = read_bpm(_W["bpm"])
        path = Path(_W["egt_dir"]) / name
        if not path.exists():
            raise FileNotFoundError(f"cluster file {name!r} named in the GTC is not in {_W['egt_dir']}")
        m = _W["manifest"]
        _W["ctx"][name] = manifest_context(m, align_clusters(read_egt(path), m), cluster_file_name=name)
    return _W["ctx"][name]


def _canon(a: np.ndarray) -> np.ndarray:  # type: ignore[type-arg]
    a = np.asarray(a).copy()
    if a.dtype.kind == "f":
        a[np.isnan(a)] = NEG_NAN
    return a


def check(args: tuple[str, str, str]) -> dict[str, Any]:
    from simurg.formats.gtc import encode_gtc, read_gtc
    from simurg.formats.idat import read_idat
    from simurg.genotype.sample import genotype_sample, params_for_preset
    from simurg.genotype.sex import SexParams

    gtc, grn, red = args
    res: dict[str, Any] = {"gtc": gtc}
    try:
        ref = read_gtc(gtc, check_write_complete=False)
        ctx = _ctx(ref.cluster_file)
        params = (
            params_for_preset(_W["preset"], gtc_version=ref.version)
            if _W["preset"]
            else params_for_preset(
                "current", norm_version=_W["norm"], sex=SexParams(version=_W["sex_version"]), gtc_version=ref.version
            )
        )
        g = genotype_sample(read_idat(grn), read_idat(red), ctx, params)
        res["cluster_file"] = ctx.cluster_file_name
        diffs: dict[str, Any] = {}
        for k in ("raw_x", "raw_y", "genotypes", "genotype_scores", "base_calls", "controls_x", "controls_y"):
            a, b = _canon(getattr(g, k)), _canon(getattr(ref, k))
            if a.shape != b.shape:
                diffs[k] = f"shape {a.shape} vs {b.shape}"
            elif a.tobytes() != b.tobytes():
                n = int((a.reshape(len(a), -1) != b.reshape(len(b), -1)).any(axis=1).sum())
                diffs[k] = n
        if g.normalization_transforms.tobytes() != ref.normalization_transforms.tobytes():
            diffs["normalization_transforms"] = True
        if ref.version >= 4:
            for k in ("b_allele_freqs", "logr_ratios"):
                if _canon(getattr(g, k)).tobytes() != _canon(getattr(ref, k)).tobytes():
                    diffs[k] = int(
                        (_canon(getattr(g, k)).view(np.uint32) != _canon(getattr(ref, k)).view(np.uint32)).sum()
                    )
        for k in (
            "call_rate",
            "gender",
            "logr_dev",
            "gc10",
            "gc50",
            "num_calls",
            "num_no_calls",
            "num_intensity_only",
            "ploidy",
            "ploidy_type",
            "dx",
        ):
            a, b = getattr(g, k), getattr(ref, k)
            same = a == b or (isinstance(a, float) and np.isnan(a) and np.isnan(b))
            if not same:
                diffs[k] = [repr(a), repr(b)]
        if ref.version >= 5 and g.slide_identifier != ref.slide_identifier:  # field 1016 exists only in v5
            diffs["slide_identifier"] = [g.slide_identifier, ref.slide_identifier]
        if "genotype_scores" in diffs:
            d = np.abs(g.genotype_scores.astype(np.float64) - ref.genotype_scores.astype(np.float64))
            res["max_score_abs_diff"] = float(np.nanmax(d))
        res["field_diffs"] = diffs
        # also evaluate the other sex-estimate generation on the same calls (cheap), to identify the vendor's rule
        from simurg.genotype.sex import estimate_sex
        from simurg.signal.normalization.apply import normalize, polar

        nx, ny = normalize(g.raw_x, g.raw_y, ctx.lookup, g.normalization_transforms)
        _, r_all = polar(nx, ny)
        res["sex_by_version"] = {
            v: estimate_sex(
                g.genotypes, r_all, ctx.clusters.gentrain, ctx.clusters.present, ctx.sex_classes, SexParams(version=v)
            )
            for v in (1, 2)
        }
        res["sex_ref"] = ref.gender_char
        for k in ALIGN:
            setattr(g, k, getattr(ref, k))
        for k in ("b_allele_freqs", "logr_ratios", "genotype_scores"):
            setattr(g, k, _canon(getattr(g, k)))
            setattr(ref, k, _canon(getattr(ref, k)))
        res["byte_identical"] = encode_gtc(g) == encode_gtc(ref)
        res["ref_layout_canonical"] = encode_gtc(ref) == Path(gtc).read_bytes() or None
        res["metadata"] = {k: str(getattr(ref, k))[:80] for k in ("autocall_version", "cluster_file", "imaging_date")}
    except Exception as e:
        res["error"] = f"{type(e).__name__}: {e}"
        res["trace"] = traceback.format_exc(limit=4)
    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-dir", required=True, type=Path)
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--norm-version", default="1.1.2")
    ap.add_argument("--sex-version", type=int, default=1)
    ap.add_argument("--preset", default="autocall-1.6.2", help="genotype preset; '' = use --norm/--sex-version")
    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.discovery import walk_files

    t0 = time.time()
    files = walk_files(a.root, (".gtc", ".idat"))
    key_re = re.compile(r"(\d+_r\d{2}c\d{2})", re.IGNORECASE)
    grn, red, gtcs = {}, {}, []
    for p in files:
        low = p.name.lower()
        m = key_re.search(p.name)
        if not m:
            continue
        k = m.group(1).lower()
        if low.endswith("_grn.idat"):
            grn[k] = str(p)
        elif low.endswith("_red.idat"):
            red[k] = str(p)
        elif low.endswith(".gtc"):
            gtcs.append((k, str(p)))
    triples = [(g, grn[k], red[k]) for k, g in gtcs if k in grn and k in red]
    if a.limit:
        triples = triples[: a.limit]
    with ProcessPoolExecutor(
        a.workers,
        mp_context=mp.get_context("spawn"),
        initializer=_init,
        initargs=(str(a.bpm), str(a.egt_dir), a.norm_version, a.sex_version, a.preset),
    ) as ex:
        results = list(ex.map(check, triples, chunksize=2))

    field_counts: dict[str, int] = {}
    for r in results:
        for k in r.get("field_diffs", {}):
            field_counts[k] = field_counts.get(k, 0) + 1
    summary: dict[str, Any] = {
        "n_gtc": len(gtcs),
        "n_compared": len(results),
        "errors": [r for r in results if "error" in r][:20],
        "byte_identical": sum(bool(r.get("byte_identical")) for r in results),
        "samples_with_no_field_diff": sum(not r.get("field_diffs") and "error" not in r for r in results),
        "samples_with_field_diff": field_counts,
        "genotype_diff_loci_total": sum(
            int(r.get("field_diffs", {}).get("genotypes", 0) or 0)
            if isinstance(r.get("field_diffs", {}).get("genotypes", 0), int)
            else 0
            for r in results
        ),
        "score_diff_loci_total": sum(
            int(r.get("field_diffs", {}).get("genotype_scores", 0) or 0)
            if isinstance(r.get("field_diffs", {}).get("genotype_scores", 0), int)
            else 0
            for r in results
        ),
        "max_score_abs_diff": max((r.get("max_score_abs_diff", 0.0) for r in results), default=0.0),
        "cluster_files": sorted({r.get("cluster_file", "") for r in results}),
        "sex_agreement_by_version": {
            str(v): sum(r.get("sex_by_version", {}).get(v) == r.get("sex_ref") for r in results) for v in (1, 2)
        },
        "ref_sex_counts": {s: sum(r.get("sex_ref") == s for r in results) for s in ("M", "F", "U")},
        "examples": [r for r in results if r.get("field_diffs")][:15],
        "params": {"preset": a.preset, "norm_version": a.norm_version, "sex_version": a.sex_version},
        "seconds": round(time.time() - t0, 1),
    }
    summary["passed"] = summary["byte_identical"] == len(results) > 0
    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_compared",
                    "byte_identical",
                    "samples_with_no_field_diff",
                    "samples_with_field_diff",
                    "genotype_diff_loci_total",
                    "score_diff_loci_total",
                    "passed",
                    "seconds",
                )
            }
        )
    )
    return 0 if summary["passed"] else 1


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