"""Full-corpus, model-only validation shared by the independent exporters.

No native application/library is required. Every dataset row is accounted for;
bad inputs and inference failures are retained in the report, never silently dropped.
"""
from __future__ import annotations

import argparse
from concurrent.futures import ProcessPoolExecutor
from collections import defaultdict
from datetime import datetime, timezone
import gc
import hashlib
import importlib.metadata
import io
import json
import math
from pathlib import Path
import platform
import time

import numpy as np

from export_common import MODEL_ID, REVISION, ROOT, load_reference, sha256_file

DATASET = "yojul/one-shot-hip-hop-drums"
DATASET_REVISION = "07be604e729fcf8f7904bca5328d6b09b1ceed7b"
DATASET_ROWS = 19673
VARIANTS = ("fp32", "int8", "int4")
COMPARISONS = ("upstream -> fp32", "fp32 -> int8", "fp32 -> int4")


def parser(description):
    result = argparse.ArgumentParser(description=description)
    mode = result.add_mutually_exclusive_group()
    mode.add_argument("--validate-only", action="store_true", help="Validate existing graphs without exporting again")
    mode.add_argument("--skip-validation", action="store_true", help="Export with smoke checks only; skip corpus evaluation")
    result.add_argument("--dataset-cache", type=Path, help="Hugging Face cache (default: standard HF cache)")
    result.add_argument("--results-dir", type=Path, default=ROOT / "validation-results")
    result.add_argument("--work-dir", type=Path, default=ROOT / ".validation-cache")
    result.add_argument("--threads", type=int, default=2)
    result.add_argument("--workers", type=int, default=1, help="Independent CPU workers; each loads its own models")
    result.add_argument("--limit", type=int, help="Diagnostic prefix only; visibly marked PARTIAL, never a full-corpus result")
    result.add_argument("--resume", action="store_true", help="Resume only an exactly matching checkpoint")
    result.add_argument("--offline", action="store_true", help="Use cached Hub files only")
    return result


def divergence(reference, actual):
    if reference.shape != actual.shape:
        raise ValueError(f"shape mismatch: {reference.shape} != {actual.shape}")
    a, b = reference.astype(np.float64).ravel(), actual.astype(np.float64).ravel()
    if not len(a) or not np.isfinite(a).all() or not np.isfinite(b).all():
        raise ValueError("empty or non-finite output")
    denominator = float(np.linalg.norm(a) * np.linalg.norm(b))
    if denominator <= 0:
        raise ValueError("zero-norm output has undefined cosine")
    difference = a - b
    return {
        "cosine": float(np.clip(np.dot(a, b) / denominator, -1, 1)),
        "max_abs": float(np.max(np.abs(difference))),
        "rmse": float(np.sqrt(np.mean(difference * difference))),
    }


def summary(values):
    values = np.asarray(values, dtype=np.float64)
    if not len(values):
        return {key: None for key in ("mean", "min", "p05", "p95", "max")}
    return {"mean": float(values.mean()), "min": float(values.min()),
            "p05": float(np.percentile(values, 5)), "p95": float(np.percentile(values, 95)),
            "max": float(values.max())}


def passes(metrics, comparison):
    # Frozen exporter tolerances. Quantized values are observations, not a claim
    # that these permissive cosine thresholds imply perceptual quality.
    if comparison == "upstream -> fp32":
        return metrics["cosine"] >= 0.99999 and metrics["max_abs"] <= 5e-4
    return metrics["cosine"] >= (0.97 if comparison.endswith("int8") else 0.80)


def decode_audio(row):
    import soundfile as sf
    from scipy.signal import resample_poly

    encoded = row["audio"]["bytes"]
    if not encoded:
        raise ValueError("dataset row contains zero audio bytes")
    samples, rate = sf.read(io.BytesIO(encoded), dtype="float32", always_2d=True)
    source_frames, channels = samples.shape
    if not source_frames or not np.isfinite(samples).all():
        raise ValueError("empty or non-finite decoded audio")
    # Explicit shared Python preparation, independent of any native runtime.
    mono = samples[:30 * rate].mean(axis=1, dtype=np.float32)
    if rate != 16000:
        divisor = math.gcd(rate, 16000)
        mono = resample_poly(mono, 16000 // divisor, rate // divisor).astype(np.float32)
    if len(mono) <= 160:
        raise ValueError("Gemma requires more than 10 ms of audio")
    return mono, {"sample_rate": rate, "channels": channels, "source_frames": source_frames,
                  "evaluated_samples": len(mono), "truncated": source_frames > 30 * rate}


def make_inputs(processor, samples):
    inputs = processor(audio=[samples], sampling_rate=16000, return_tensors="np")
    dtypes = {"input_ids": np.int64, "attention_mask": np.int64,
              "input_features": np.float32, "input_features_mask": np.bool_}
    return {name: inputs[name].astype(dtype) for name, dtype in dtypes.items()}


def load_upstream(kind, config):
    import torch
    if kind == "tower":
        from huggingface_hub import hf_hub_download
        from safetensors import safe_open
        from transformers.models.gemma4.modeling_gemma4 import Gemma4AudioModel
        from export_tower import AudioTower

        model = Gemma4AudioModel(config.audio_config)
        checkpoint = hf_hub_download(MODEL_ID, "model.safetensors", revision=REVISION)
        with safe_open(checkpoint, framework="pt", device="cpu") as weights:
            state = {key.removeprefix("audio_tower."): weights.get_tensor(key)
                     for key in weights.keys() if key.startswith("audio_tower.")}
        model.load_state_dict(state, strict=True)
        return AudioTower(model).eval()
    from transformers import AutoModel
    from export_embedding import AudioEmbedding

    config.vision_config = None
    return AudioEmbedding(AutoModel.from_pretrained(
        MODEL_ID, revision=REVISION, config=config, dtype=torch.float32,
    ).eval()).eval()


def upstream_output(model, kind, feed):
    import torch
    tensors = {name: torch.from_numpy(value) for name, value in feed.items()}
    with torch.inference_mode():
        if kind == "tower":
            states, mask = model(tensors["input_features"], tensors["input_features_mask"])
            return states.cpu().float().numpy(), mask.cpu().numpy()
        return model(**tensors).cpu().float().numpy(), None


def ort_output(session, feed, kind):
    outputs = session.run(None, {value.name: feed[value.name] for value in session.get_inputs()})
    return outputs[0], outputs[1] if kind == "tower" else None


def output_contract(values, mask, kind):
    if not np.isfinite(values).all():
        raise ValueError("non-finite model output")
    if kind == "embedding":
        if values.shape != (1, 512) or not np.isclose(np.linalg.norm(values), 1, atol=1e-5):
            raise ValueError("embedding shape or L2 normalization contract failed")
    elif (values.ndim != 3 or values.shape[0] != 1 or values.shape[2] != 1536
          or mask is None or mask.shape != values.shape[:2] or mask.dtype != np.bool_):
        raise ValueError("tower shape/mask contract failed")


def bucket_key(shape, mask):
    # Raw towers are never pooled or compared across incompatible shapes/masks.
    return json.dumps([shape, mask], separators=(",", ":"))


def nearest_indices(vectors, k):
    """Exact cosine neighbors, self excluded; row index breaks exact ties."""
    norms = np.linalg.norm(vectors.astype(np.float64), axis=1)
    normalized = (vectors / norms[:, None]).astype(np.float32)
    output = np.empty((len(vectors), k), dtype=np.int32)
    for start in range(0, len(vectors), 128):
        scores = normalized[start:start + 128] @ normalized.T
        scores[np.arange(len(scores)), np.arange(start, start + len(scores))] = -np.inf
        output[start:start + len(scores)] = np.argsort(-scores, axis=1, kind="stable")[:, :k]
    return output


def neighbor_agreement(records, work):
    groups = defaultdict(list)
    for record in records:
        if record["status"] == "ok":
            groups[bucket_key(record["shape"], record["mask"])].append(record)
    maps = {variant: np.memmap(work / f"{variant}.f32", dtype="<f4", mode="r")
            for variant in VARIANTS} if any(groups.values()) else {}
    totals = {variant: {"eligible": 0, "top1_matches": 0, "overlap_at_10": [], "top10_eligible": 0}
              for variant in ("int8", "int4")}
    singleton = 0
    for group_number, group in enumerate(groups.values()):
        if len(group) < 2:
            singleton += len(group)
            continue
        k = min(10, len(group) - 1)
        def matrix(variant):
            return np.stack([maps[variant][r["offset"]:r["offset"] + r["count"]] for r in group])
        reference = nearest_indices(matrix("fp32"), k)
        for variant in ("int8", "int4"):
            actual = nearest_indices(matrix(variant), k)
            total = totals[variant]
            total["eligible"] += len(group)
            total["top1_matches"] += int((reference[:, 0] == actual[:, 0]).sum())
            if k == 10:
                total["top10_eligible"] += len(group)
                total["overlap_at_10"].extend(
                    len(set(a) & set(b)) / 10 for a, b in zip(reference, actual)
                )
        if group_number % 50 == 0:
            print(f"Neighbor groups: {group_number + 1}/{len(groups)}", flush=True)
    return {variant: {"eligible": total["eligible"], "singleton_excluded": singleton,
                      "top1_agreement": total["top1_matches"] / total["eligible"] if total["eligible"] else None,
                      "top10_eligible": total["top10_eligible"],
                      "mean_overlap_at_10": float(np.mean(total["overlap_at_10"])) if total["overlap_at_10"] else None}
            for variant, total in totals.items()}


SWEEPS = [("pitch", [110., 220., 440.]), ("pitch", [220., 330., 440.]),
          ("harmonic_mix", [0., .3, .9]), ("harmonic_mix", [.1, .5, 1.]),
          ("filter_cutoff", [300., 1000., 5000.]), ("filter_cutoff", [500., 2000., 8000.]),
          ("decay", [.05, .2, .8]), ("decay", [.1, .4, 1.]),
          ("fm_depth", [0., 1., 4.]), ("fm_depth", [.5, 2., 6.])]


def synth_signal(knob, value):
    """Frozen one-second, one-knob sweeps implemented entirely in Python."""
    rate = 16000
    t = np.arange(rate, dtype=np.float64) / rate
    hz = value if knob == "pitch" else 220.
    phase = 2 * np.pi * hz * t
    fm = value * np.sin(2 * np.pi * 110 * t) if knob == "fm_depth" else 0.
    harmonic = value if knob == "harmonic_mix" else .3
    raw = np.sin(phase + fm) + harmonic * np.sin(2 * phase) + harmonic * .5 * np.sin(3 * phase)
    cutoff = value if knob == "filter_cutoff" else 8000.
    alpha = 1 - math.exp(-2 * math.pi * cutoff / rate)
    low = 0.
    filtered = np.empty(rate)
    for i, sample in enumerate(raw):
        low += alpha * (sample - low)
        filtered[i] = low
    decay = value if knob == "decay" else .6
    envelope = np.minimum(t / .01, 1) * np.exp(-t / decay) * np.clip(1 - t, 0, .01) / .01
    return (.3 * filtered * envelope).astype(np.float32)


def continuity(model, sessions, processor, kind):
    results = {variant: [] for variant in ("upstream", *VARIANTS)}
    for knob, values in SWEEPS:
        outputs = {variant: [] for variant in results}
        for value in values:
            feed = make_inputs(processor, synth_signal(knob, value))
            outputs["upstream"].append(upstream_output(model, kind, feed)[0])
            for variant, session in sessions.items():
                outputs[variant].append(ort_output(session, feed, kind)[0])
        for variant, (a, b, c) in outputs.items():
            distances = [1 - divergence(a, b)["cosine"], 1 - divergence(b, c)["cosine"],
                         1 - divergence(a, c)["cosine"]]
            results[variant].append({"knob": knob, "values": values, "distances": distances,
                                     "passed": max(distances[:2]) <= distances[2] + 1e-6})
    return {variant: {"passed": sum(case["passed"] for case in cases), "total": len(cases),
                      "required_fraction": .9, "expectation_met": sum(c["passed"] for c in cases) >= 9,
                      "cases": cases} for variant, cases in results.items()}


def markdown(report):
    def number(value):
        return "N/A" if value is None else f"{value:.6g}"
    lines = [f"### {report['kind'].title()} — {'FULL CORPUS' if report['full_corpus'] else 'PARTIAL'}",
             "", f"Dataset `{DATASET}` at `{DATASET_REVISION}`.",
             f"Attempted **{report['attempted']:,}/{report['dataset_rows']:,}** rows; "
             f"successful **{report['successful']:,}**; input errors **{report['input_errors']:,}**; "
             f"inference errors **{report['inference_errors']:,}**.", "",
             "| Comparison | Clips | Within tolerance | Mean cosine | Min cosine | P95 clip max abs | Worst abs | P95 RMSE |",
             "| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: |"]
    for name, row in report["divergence"].items():
        lines.append(f"| {name} | {row['count']:,} | {row['passed']:,}/{row['count']:,} | "
                     f"{number(row['cosine']['mean'])} | {number(row['cosine']['min'])} | "
                     f"{number(row['max_abs']['p95'])} | {number(row['max_abs']['max'])} | {number(row['rmse']['p95'])} |")
    lines += ["", "| Quantization vs FP32 | Top-1 eligible queries | Top-1 agreement | Top-10 eligible queries | Mean overlap@10 |",
              "| --- | ---: | ---: | ---: | ---: |"]
    for variant, row in report["neighbors"].items():
        lines.append(f"| {variant.upper()} | {row['eligible']:,} | {number(row['top1_agreement'])} | "
                     f"{row['top10_eligible']:,} | {number(row['mean_overlap_at_10'])} |")
    lines += ["", "| Synth continuity | Passing triples | Meets 90% expectation |", "| --- | ---: | --- |"]
    for variant, row in report["continuity"].items():
        lines.append(f"| {variant} | {row['passed']}/{row['total']} | {'Yes' if row['expectation_met'] else 'No'} |")
    lines += ["", "Raw-tower neighbor searches use only identical shapes and masks; no pooling is applied. "
              "Final embeddings search all successful clips. Self matches are excluded; exact ties use dataset row order. "
              "Overlap@10 includes only groups with at least 11 clips.",
              "", "FP32 tolerance: cosine >= 0.99999 and max absolute error <= 0.0005. "
              "INT8/INT4 cosine tolerances are 0.97/0.80, matching the existing smoke gates; "
              "they are numerical diagnostics, not perceptual acceptance criteria.",
              "", "Python-versus-native preprocessing parity is **deferred** and was not measured here.", ""]
    return "\n".join(lines)


_CONTEXT = None


def initialize_worker(kind, threads, offline):
    """Process-local pristine reference and ONNX sessions; no export patches."""
    import onnxruntime as ort
    import torch
    if offline:
        import os
        import huggingface_hub.constants
        os.environ["HF_HUB_OFFLINE"] = "1"
        huggingface_hub.constants.HF_HUB_OFFLINE = True
    graph_names = {v: f"audio_{kind}{'' if v == 'fp32' else '_' + v}.onnx" for v in VARIANTS}
    config, processor = load_reference()
    torch.set_num_threads(threads)
    upstream = load_upstream(kind, config)
    sessions = {}
    for variant, name in graph_names.items():
        options = ort.SessionOptions()
        options.intra_op_num_threads = threads
        options.inter_op_num_threads = 1
        options.log_severity_level = 3
        sessions[variant] = ort.InferenceSession(str(ROOT / name), options, providers=["CPUExecutionProvider"])
    global _CONTEXT
    _CONTEXT = kind, processor, upstream, sessions


def evaluate_case(item):
    row_number, row = item
    kind, processor, upstream, sessions = _CONTEXT
    record = {"row": row_number, "path": row["audio"]["path"], "label": row["label"]}
    try:
        samples, preparation = decode_audio(row)
        feed = make_inputs(processor, samples)
        record.update(preparation)
    except Exception as error:
        record.update(status="input_error", error=f"{type(error).__name__}: {error}")
        return record, None
    try:
        reference, ref_mask = upstream_output(upstream, kind, feed)
        output_contract(reference, ref_mask, kind)
        outputs = {}
        for variant, session in sessions.items():
            values, mask = ort_output(session, feed, kind)
            output_contract(values, mask, kind)
            if not np.array_equal(mask, ref_mask):
                raise ValueError(f"{variant} mask differs from upstream")
            outputs[variant] = values
        metrics = {"upstream -> fp32": divergence(reference, outputs["fp32"]),
                   "fp32 -> int8": divergence(outputs["fp32"], outputs["int8"]),
                   "fp32 -> int4": divergence(outputs["fp32"], outputs["int4"])}
        record.update(status="ok", metrics=metrics, count=outputs["fp32"].size,
                      shape=list(reference.shape), mask=ref_mask.ravel().tolist() if ref_mask is not None else [])
        return record, outputs
    except Exception as error:
        record.update(status="inference_error", error=f"{type(error).__name__}: {error}")
        return record, None


def run_validation(kind, args):
    global _CONTEXT
    import onnxruntime as ort
    import pyarrow.parquet as pq
    import torch
    from huggingface_hub import snapshot_download

    if args.threads < 1 or args.workers < 1 or (args.limit is not None and args.limit < 1):
        raise ValueError("threads and limit must be positive")
    if args.offline:
        # Set the library's constant too: huggingface_hub may have been imported earlier.
        import os
        import huggingface_hub.constants
        os.environ["HF_HUB_OFFLINE"] = "1"
        huggingface_hub.constants.HF_HUB_OFFLINE = True
    snapshot = Path(snapshot_download(DATASET, repo_type="dataset", revision=DATASET_REVISION,
                                     cache_dir=args.dataset_cache, local_files_only=args.offline))
    shards = sorted((snapshot / "data").glob("*.parquet"))
    total = sum(pq.ParquetFile(path).metadata.num_rows for path in shards)
    if total != DATASET_ROWS:
        raise ValueError(f"Pinned corpus row count changed: {total} != {DATASET_ROWS}")
    target = min(args.limit, total) if args.limit else total
    graph_names = {v: f"audio_{kind}{'' if v == 'fp32' else '_' + v}.onnx" for v in VARIANTS}
    hashes = {name: sha256_file(ROOT / name) for name in graph_names.values()}
    if kind == "embedding":
        hashes["audio_embedding.data"] = sha256_file(ROOT / "audio_embedding.data")
    versions = {name: importlib.metadata.version(name) for name in
                ("torch", "transformers", "onnxruntime", "numpy", "scipy", "soundfile", "pyarrow", "huggingface_hub")}
    code = {name: sha256_file(ROOT / name) for name in ("validation.py", "export_common.py", f"export_{kind}.py")}
    provenance = {"kind": kind, "dataset": DATASET, "dataset_revision": DATASET_REVISION,
                  "dataset_files": {p.name: sha256_file(p) for p in shards}, "dataset_rows": total,
                  "target": target, "model": MODEL_ID, "model_revision": REVISION,
                  "artifacts": hashes, "versions": versions, "code": code, "threads": args.threads, "workers": args.workers,
                  "python": platform.python_version(), "platform": platform.platform()}
    identity = hashlib.sha256(json.dumps(provenance, sort_keys=True).encode()).hexdigest()
    work = args.work_dir / f"{kind}-{identity[:16]}"
    work.mkdir(parents=True, exist_ok=True)
    log_path = work / "cases.jsonl"
    if log_path.exists() and not args.resume:
        raise ValueError(f"Checkpoint already exists; use --resume or another --work-dir: {work}")
    records = []
    if log_path.exists():
        with log_path.open("rb") as stream:
            good_end = 0
            for line in stream:
                try:
                    records.append(json.loads(line))
                except json.JSONDecodeError:
                    break  # A interrupted final write is not a completed row.
                good_end = stream.tell()
        with log_path.open("r+b") as stream:
            stream.truncate(good_end)
        if [r["row"] for r in records] != list(range(len(records))):
            raise ValueError("Checkpoint row sequence is corrupt")
    offset = sum(r.get("count", 0) for r in records if r["status"] == "ok")
    streams = {}
    for variant in VARIANTS:
        path = work / f"{variant}.f32"
        stream = path.open("r+b" if path.exists() else "w+b")
        if path.stat().st_size < offset * 4:
            raise ValueError("Checkpoint vectors are incomplete")
        stream.truncate(offset * 4)
        stream.seek(0, 2)
        streams[variant] = stream
    print(f"EXPERIMENTAL {kind}: {len(records):,}/{target:,} rows checkpointed; cache {work}", flush=True)
    started, resumed = time.monotonic(), len(records)
    def remaining_rows():
        row_number = -1
        for shard in shards:
            for batch in pq.ParquetFile(shard).iter_batches(batch_size=128):
                for row in batch.to_pylist():
                    row_number += 1
                    if row_number >= target:
                        return
                    if row_number >= resumed:
                        yield row_number, row
    pool = None
    try:
        if args.workers == 1:
            initialize_worker(kind, args.threads, args.offline)
            results = map(evaluate_case, remaining_rows())
        else:
            pool = ProcessPoolExecutor(max_workers=args.workers, initializer=initialize_worker,
                                       initargs=(kind, args.threads, args.offline))
            results = pool.map(evaluate_case, remaining_rows(), chunksize=4)
        with log_path.open("a", encoding="utf-8") as log:
            for record, outputs in results:
                if record["status"] == "ok":
                    record["offset"] = offset
                    for variant in VARIANTS:
                        streams[variant].write(outputs[variant].astype("<f4").tobytes())
                        streams[variant].flush()
                    offset += record["count"]
                log.write(json.dumps(record, allow_nan=False) + "\n")
                log.flush()
                records.append(record)
                completed = len(records) - resumed
                if completed == 1 or len(records) % 50 == 0 or len(records) == target:
                    elapsed = time.monotonic() - started
                    speed = completed / max(elapsed, .001)
                    print(f"{kind}: {len(records):,}/{target:,}; {speed:.2f} rows/s; "
                          f"ETA {(target - len(records)) / max(speed, .001) / 60:.1f} min", flush=True)
    finally:
        if pool is not None:
            pool.shutdown(wait=True, cancel_futures=True)
        for stream in streams.values():
            stream.close()
    if args.workers != 1:
        initialize_worker(kind, args.threads, args.offline)
    _, processor, upstream, sessions = _CONTEXT
    synth = continuity(upstream, sessions, processor, kind)
    _CONTEXT = None
    del upstream, sessions
    gc.collect()
    neighbors = neighbor_agreement(records, work)
    successful = [r for r in records if r["status"] == "ok"]
    report = dict(provenance, generated_utc=datetime.now(timezone.utc).isoformat(), full_corpus=target == total,
                  attempted=len(records), successful=len(successful),
                  input_errors=sum(r["status"] == "input_error" for r in records),
                  inference_errors=sum(r["status"] == "inference_error" for r in records),
                  truncated=sum(r.get("truncated", False) for r in records),
                  divergence={}, neighbors=neighbors, continuity=synth,
                  errors=[r for r in records if r["status"] != "ok"],
                  deferred=["Python-versus-native preprocessing parity"], checkpoint_identity=identity)
    for comparison in COMPARISONS:
        values = [r["metrics"][comparison] for r in successful]
        report["divergence"][comparison] = {"count": len(values),
            "passed": sum(passes(v, comparison) for v in values),
            **{metric: summary([v[metric] for v in values]) for metric in ("cosine", "max_abs", "rmse")}}
    args.results_dir.mkdir(parents=True, exist_ok=True)
    label = kind if target == total else f"{kind}-partial-{target}"
    (args.results_dir / f"{label}.json").write_text(json.dumps(report, indent=2, allow_nan=False) + "\n", encoding="utf-8")
    table = markdown(report)
    (args.results_dir / f"{label}.md").write_text(table, encoding="utf-8")
    print("\n" + table, flush=True)
    print(f"Machine-readable report: {args.results_dir / (label + '.json')}", flush=True)
    # Record all observations before returning a failing gate; never hide failures.
    return (not report["input_errors"] and not report["inference_errors"]
            and all(r["passed"] == r["count"] for r in report["divergence"].values())
            and all(r["expectation_met"] for r in synth.values()))