#!/usr/bin/env python3
"""Reproducible CPU validation and sequential CUDA experiments (no network)."""
from __future__ import annotations

import argparse
import hashlib
import json
import platform
import random
import statistics
import subprocess
import time
from datetime import datetime
from pathlib import Path
from zoneinfo import ZoneInfo

HERE = Path(__file__).resolve().parent
POS = HERE / "positions"
OUT = HERE / "results"
BIN = "/tmp/slime_feasibility_cuda"


def save_positions(name, points):
    (POS / f"{name}.txt").write_text("".join(f"{x} {z}\n" for x, z in points))


def prepare():
    POS.mkdir(exist_ok=True)
    OUT.mkdir(exist_ok=True)
    save_positions("known", [(51069, 419124)])
    offsets = [(0, 0), (0, 1), (1, 0), (-1, 0), (0, -1), (2, 0), (0, 2),
               (2, 2), (-2, -2), (-1, 1), (1, -1), (3, 0), (0, 3), (3, 3),
               (-3, 0), (0, -3)]
    save_positions("stress", [(51069 + x, 419124 + z) for x, z in offsets] +
                   [(0, 0), (-1875000, -1875000), (1874995, 1874995),
                    (2147483643, 2147483643), (-2147483648, -2147483648)])
    rng = random.Random(2026101101)
    coords = [(rng.randrange(-1875000, 1875000), rng.randrange(-1875000, 1875000))
              for _ in range(65536)]
    save_positions("random4096", coords[:4096])
    save_positions("random65536", coords)
    save_positions("tail128", [(51069, 419124)] + coords[:127])
    return coords


def run_case(name, positions, task, repeats=5, cells=None):
    cmd = [BIN, str(POS / f"{positions}.txt"), task, str(repeats)]
    if cells is not None:
        cmd.append(str(cells))
    start = time.perf_counter()
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
    (OUT / f"{name}.jsonl").write_text(result.stdout)
    (OUT / f"{name}.stderr").write_text(result.stderr)
    if result.returncode:
        raise RuntimeError(f"{name} failed: {result.stderr}")
    rows = [json.loads(line) for line in result.stdout.splitlines()]
    record = {"command": cmd, "wall_ms": (time.perf_counter() - start) * 1000,
              "rows": rows}
    print(json.dumps({"finished": name, "wall_ms": record["wall_ms"]}), flush=True)
    return record


def median(rows, kind, mode):
    return statistics.median(r["kernel_ms"] for r in rows
                             if r["type"] == kind and r["mode"] == mode)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--analyse-only", action="store_true")
    ap.add_argument("--complete-only", action="store_true")
    args = ap.parse_args()
    if args.analyse_only or args.complete_only:
        records = json.loads((OUT / "all.json").read_text())
    else:
        coords = prepare()
        subprocess.run(["nvcc", "-O3", "-std=c++17", "-arch=sm_86",
                        str(HERE / "cuda_bench.cu"), "-o", BIN], check=True)
        subprocess.run(["g++", "-O3", "-std=c++17", "-march=native", "-Wall", "-Wextra",
                        str(HERE / "cpu_verify.cpp"), "-o", "/tmp/slime_feasibility_verify"], check=True)
        cpu = subprocess.check_output(["/tmp/slime_feasibility_verify"], text=True)
        (OUT / "cpu_verify.json").write_text(cpu)
        gpu_query = ["nvidia-smi", "--query-gpu=name,driver_version,temperature.gpu,utilization.gpu", "--format=csv"]
        records = {"metadata": {
            "date": datetime.now(ZoneInfo("Asia/Hong_Kong")).isoformat(),
            "platform": platform.platform(), "cpu_validation": json.loads(cpu),
            "gpu_before": subprocess.check_output(gpu_query, text=True),
            "source_sha256": {p.name: hashlib.sha256(p.read_bytes()).hexdigest()
                              for p in [HERE / "cuda_bench.cu", HERE / "interval_solver.hpp",
                                        HERE / "cpu_verify.cpp", HERE / "run.py",
                                        HERE.parent.parent / "slime_solver_core.hpp",
                                        HERE.parents[2] / "include/slime/solver_core.hpp",
                                        HERE.parents[2] / "include/slime/upper_interval.hpp"]}}}

        def record(name, positions, task, repeats=5, cells=None):
            records[name] = run_case(name, positions, task, repeats, cells)
            (OUT / "all.json").write_text(json.dumps(records, indent=2))

        record("known", "known", "compare")
        record("stress", "stress", "compare")
        record("random4096", "random4096", "compare")
        record("random65536", "random65536", "fast")
        record("roots65536", "random65536", "root")
        counts = next(r["counts"] for r in records["random4096"]["rows"] if r["type"] == "root_counts")
        classes = [(x % 65536, z % 131072) for (x, z), n in zip(coords, counts) if n][:32]
        save_positions("classes32", classes)
        lifts, groups = [], []
        for x, z in classes:
            current = [(xx, zz)
                       for xx in range(-1875000 + (x + 1875000) % 65536, 1875000, 65536)
                       for zz in range(-1875000 + (z + 1875000) % 131072, 1875000, 131072)]
            groups.append(len(current))
            lifts.extend(current)
        save_positions("lifts", lifts)
        (POS / "lift_groups.json").write_text(json.dumps(groups))
        record("classes32", "classes32", "root")
        record("lifts", "lifts", "root")
        record("tail_one_cell", "known", "tail", cells=1)
        record("tail_two_cells", "known", "tail", cells=2)
        record("tail_five_cells", "known", "tail", cells=5)
        record("tail128", "tail128", "tail")
        records["metadata"]["gpu_after"] = subprocess.check_output(gpu_query, text=True)
        (OUT / "all.json").write_text(json.dumps(records, indent=2))

    if not args.analyse_only:
        if args.complete_only:
            subprocess.run(["nvcc", "-O3", "-std=c++17", "-arch=sm_86",
                            str(HERE / "cuda_bench.cu"), "-o", BIN], check=True)
        records["complete_stress"] = run_case("complete_stress", "stress", "complete")
        records["complete65536"] = run_case("complete65536", "random65536", "complete")
        records["metadata"]["complete_source_sha256"] = {
            p.name: hashlib.sha256(p.read_bytes()).hexdigest()
            for p in [HERE / "cuda_bench.cu", HERE / "interval_solver.hpp", HERE / "run.py",
                      HERE.parents[2] / "include/slime/solver_core.hpp",
                      HERE.parents[2] / "include/slime/upper_interval.hpp"]}
        (OUT / "all.json").write_text(json.dumps(records, indent=2))

    summary = {}
    for name in ["known", "stress", "random4096", "random65536"]:
        rows = records[name]["rows"]
        modes = sorted({r["mode"] for r in rows if r["type"] == "search_timing"})
        summary[name] = {"kernel_median_ms": {str(m): median(rows, "search_timing", m) for m in modes}}
        summary[name].update(next(r for r in rows if r["type"] == "search_info"))
    counts = next(r["counts"] for r in records["random65536"]["rows"] if r["type"] == "root_counts")
    n = len(counts)
    mean = statistics.mean(counts)
    sem = statistics.stdev(counts) / n**0.5
    summary["distribution"] = {"positions": n, "roots": sum(counts), "mean": mean,
                               "mean_normal95": [mean - 1.96 * sem, mean + 1.96 * sem],
                               "zero_fraction": counts.count(0) / n, "max": max(counts)}
    for name in ["roots65536", "classes32", "lifts"]:
        summary[name] = {str(m): median(records[name]["rows"], "root_timing", m) for m in [0, 1]}
    for name in ["tail_one_cell", "tail_two_cells", "tail_five_cells", "tail128"]:
        rows = records[name]["rows"]
        summary[name] = {"kernel_median_ms": {str(m): median(rows, "tail_timing", m) for m in [0, 1]},
                         "matches": next(r["matches"] for r in rows if r["type"] == "tail_timing")}

    # All projections are throughput extrapolations, not measured full-world scans.
    world = 3750000**2
    year = 365.25 * 86400
    fast_ms = summary["random65536"]["kernel_median_ms"]["3"]
    tail_ms = summary["tail128"]["kernel_median_ms"]["1"]
    root_ms = min(summary["roots65536"].values())
    summary["projections"] = {
        "world_anchors": world,
        "lower35_visits": world * mean * 131072,
        "first_draw_kernel_years": fast_ms / 1000 / n * world / year,
        "root_no_reuse_years": root_ms / 1000 / n * world / year,
        "root_class_reuse_hours": root_ms / 1000 / n * (2**33) / 3600,
        "canonical_tail_kernel_years": tail_ms / 1000 / 128 * world / year,
        "complete_compute_years": (fast_ms / 1000 / n + tail_ms / 1000 / 128) * world / year,
    }
    for name in ["complete_stress", "complete65536"]:
        if name in records:
            rows = records[name]["rows"]
            summary[name] = {key: statistics.median(r[key] for r in rows if r["type"] == "complete_timing")
                             for key in ["first_ms", "tail_ms", "compute_ms", "scan_ms"]}
            summary[name].update(next(r for r in rows if r["type"] == "complete_info"))
    if "complete65536" in summary:
        summary["projections"]["complete_same_batch_kernel_years"] = summary["complete65536"]["compute_ms"] / 1000 / n * world / year
        summary["projections"]["complete_same_batch_scan_years"] = summary["complete65536"]["scan_ms"] / 1000 / n * world / year
        assert summary["complete_stress"]["matches"] == summary["stress"]["matches"]
    # Count equality is checked here; the CUDA program additionally compares complete
    # normalized root lists for every repeated coordinate class in its input.
    base_counts = next(r["counts"] for r in records["classes32"]["rows"] if r["type"] == "root_counts")
    lifted_counts = next(r["counts"] for r in records["lifts"]["rows"] if r["type"] == "root_counts")
    groups = json.loads((POS / "lift_groups.json").read_text())
    assert lifted_counts == [c for c, number in zip(base_counts, groups) for _ in range(number)]
    summary["cache_counts_equal"] = True
    (OUT / "summary.json").write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()