#!/usr/bin/env python3
"""Run sequential native KNN benchmarks and preserve complete raw evidence.

Build the same tests/knn_benchmark.cpp twice: once against original ml.cpp with
ANER_KNN_BASELINE defined, once against the updated library. Compilation and
other tests must finish before running this driver.
"""

import argparse
import hashlib
import json
import os
from pathlib import Path
import platform
import subprocess
import sys
import time


def identity(path):
    path = Path(path).resolve()
    return {"path": str(path), "sha256": hashlib.sha256(path.read_bytes()).hexdigest()}


def source_identity(root):
    if not root:
        return None
    root = Path(root).resolve()
    return {str(relative): identity(root / relative) for relative in (
        Path("src/ml.cpp"), Path("include/aner/ml.hpp"),
        Path("src/parallel.cpp"), Path("CMakeLists.txt"),
    ) if (root / relative).is_file()}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--baseline", required=True, type=Path)
    parser.add_argument("--optimized", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--source-root", type=Path)
    parser.add_argument("--baseline-source-root", type=Path)
    parser.add_argument("--n", type=int, default=2048)
    parser.add_argument("--q", type=int, default=256)
    parser.add_argument("--d", type=int, default=16)
    parser.add_argument("--ks", nargs="+", type=int, default=[5, 32])
    parser.add_argument("--workers", nargs="+", type=int, default=[1, 2, 4])
    parser.add_argument("--repeats", type=int, default=5)
    parser.add_argument("--warmups", type=int, default=1)
    parser.add_argument("--seed", type=int, default=20261030)
    args = parser.parse_args()
    if args.output.exists():
        parser.error("output already exists; preserve previous raw measurements")
    cpu_model = platform.processor()
    if platform.system() == "Darwin":
        observed = subprocess.run(["/usr/sbin/sysctl", "-n", "machdep.cpu.brand_string"],
                                  capture_output=True, text=True)
        if observed.returncode == 0:
            cpu_model = observed.stdout.strip()
    evidence = {
        "schema": "aner.knn-native-benchmark/1",
        "timestamp_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "machine": {
            "platform": platform.platform(), "system": platform.system(),
            "release": platform.release(), "machine": platform.machine(),
            "processor": platform.processor(), "cpu_model": cpu_model, "logical_cpu_count": os.cpu_count(),
        },
        "binaries": {"baseline": identity(args.baseline), "optimized": identity(args.optimized)},
        "sources": source_identity(args.source_root),
        "baseline_sources": source_identity(args.baseline_source_root),
        "driver": identity(__file__),
        "benchmark_source": identity(Path(__file__).resolve().parent.parent / "tests/knn_benchmark.cpp"),
        "notes": [
            "Cases execute sequentially; do not run builds or tests concurrently.",
            "Original baseline links the preserved pre-change native ml.cpp and uses uniform serial KNN.",
            "Timed regions measure predict() after fit; cold cases include pool creation/destruction.",
            "Complete subprocess wall time also includes independent correctness verification and is not the kernel timer.",
            "Every raw sample verifies exact predictions after stopping its timer; all readouts are checked before timing.",
            "Distance weighting is an optional changed classifier and has no pre-change weighted speed baseline.",
            "CPU slots and worker counts do not measure instantaneous hardware utilization; RSS is not measured.",
        ],
        "records": [], "skipped": [],
    }

    def run(binary, k, weights="uniform", mode="serial", workers=1):
        command = [str(binary.resolve()), "--n", str(args.n), "--q", str(args.q), "--d", str(args.d),
                   "--k", str(k), "--weights", weights, "--mode", mode, "--workers", str(workers),
                   "--repeats", str(args.repeats), "--warmups", str(args.warmups), "--seed", str(args.seed)]
        started = time.monotonic()
        completed = subprocess.run(command, capture_output=True, text=True, check=True, timeout=300)
        elapsed = time.monotonic() - started
        record = json.loads(completed.stdout)
        expected_flavor = "original" if binary == args.baseline else "optimized"
        if record.get("flavor") != expected_flavor:
            raise RuntimeError(f"{binary} is not the required {expected_flavor} build")
        record["command"] = command
        record["subprocess_wall_seconds"] = elapsed
        if completed.stderr:
            record["stderr"] = completed.stderr
        evidence["records"].append(record)
        print(f"{record['flavor']} k={k} {weights} {mode}/{workers}: "
              f"{record['median_seconds']:.6f}s median; MAD={record['mad_seconds']:.6f}s", file=sys.stderr)
        return record

    for k in args.ks:
        baseline = run(args.baseline, k)
        optimized = run(args.optimized, k)
        optimized["speedup_vs_original_serial"] = baseline["median_seconds"] / optimized["median_seconds"]
        for weights in ("uniform", "distance"):
            serial = optimized if weights == "uniform" else run(args.optimized, k, weights)
            for workers in args.workers:
                if workers > serial["cpu_hard_limit"]:
                    evidence["skipped"].append({"k": k, "weights": weights, "workers": workers,
                                                "reason": "exceeds discovered CPU allocation"})
                    continue
                for mode in ("context", "cold"):
                    record = run(args.optimized, k, weights, mode, workers)
                    record["speedup_vs_optimized_serial"] = serial["median_seconds"] / record["median_seconds"]
                    if weights == "uniform":
                        record["speedup_vs_original_serial"] = baseline["median_seconds"] / record["median_seconds"]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("x", encoding="utf-8") as output:
        json.dump(evidence, output, indent=2)
        output.write("\n")
    print(str(args.output.resolve()))


if __name__ == "__main__":
    main()