"""Float64 full-graph numerical audit using sealed Phase 4AE CUDA operators.

This synchronous diagnostic compares every neuron at every tick with the CPU
memmaps. It does not introduce a new integration rule, APL, training or tuning.
"""
from functools import lru_cache
import csv
import hashlib
import json
import math
from pathlib import Path
import sys
import time
import traceback
import types

import numpy as np
import torch

from context import ROOT
from oracle_full import queue_summary


NEURONS = 166700
EDGES = 25582938
COEFFICIENT_NAMES = ("gd", "c0", "cg", "cv", "threshold", "reset", "seed_volt")
REQUEST_FIELDS = ("tick", "neuron_idx", "attempted", "accepted", "rejected")
AE_MODEL = ROOT / "models/male-cns-lif/author-runtime-lab-v1"
ALLOWED_CASES = {(100, ()), (500, (1381,)), (500, (1197, 1381, 1511, 82312, 133666))}


def require(condition, message):
    if not condition:
        raise ValueError(message)


def write_json(path, value):
    with Path(path).open("x", encoding="utf-8") as output:
        json.dump(value, output, indent=2, allow_nan=False)
        output.write("\n")


def digest(path):
    value = hashlib.sha256()
    with Path(path).open("rb") as source:
        for block in iter(lambda: source.read(8 << 20), b""):
            value.update(block)
    return value.hexdigest()


def array_digest(array, dtype):
    return hashlib.sha256(np.ascontiguousarray(array, dtype=dtype).tobytes()).hexdigest()


def bitwise_comparison(actual, expected, floating=False):
    require(actual.shape == expected.shape and actual.dtype == expected.dtype, "Comparison shape/dtype mismatch")
    aa = actual.view(np.uint64) if floating else actual
    bb = expected.view(np.uint64) if floating else expected
    mismatch = aa != bb
    count = int(np.count_nonzero(mismatch))
    first = int(np.flatnonzero(mismatch)[0]) if count else None
    result = {"bitwise_exact": count == 0, "mismatching_neurons": count,
              "first_neuron_idx": first}
    if floating:
        finite = bool(np.isfinite(actual).all() and np.isfinite(expected).all())
        result.update(finite=finite, max_abs_volt=float(np.max(np.abs(actual - expected))) if finite else None)
        if first is not None:
            result.update(actual_bits_hex=f"{int(aa[first]):016x}", expected_bits_hex=f"{int(bb[first]):016x}")
    return result


@lru_cache(maxsize=1)
def sealed_backend():
    """Load reviewed source without writing bytecode into the sealed model."""
    source = AE_MODEL / "gpu_backend.py"
    module = types.ModuleType("phase4af_sealed_4ae_gpu_backend")
    module.__file__ = str(source)
    exec(compile(source.read_bytes(), str(source), "exec"), module.__dict__)
    return module


@lru_cache(maxsize=2)
def load_graph(graph_directory):
    graph_dir = Path(graph_directory)
    require(graph_dir.is_absolute(), "graph_dir must be absolute")
    dtypes = {"nodes": np.int32, "edge_id": np.int64, "src": np.int32, "dst": np.int32,
              "signed_raw": np.int64, "ptr": np.int64}
    graph = {}
    for name, dtype in dtypes.items():
        value = np.load(graph_dir / (name + ".npy"), mmap_mode="r", allow_pickle=False)
        require(value.dtype == dtype and value.ndim == 1, "Invalid graph array: " + name)
        graph[name] = value
    require(graph["nodes"].shape == (NEURONS,) and
            np.array_equal(graph["nodes"], np.arange(NEURONS, dtype=np.int32)), "Graph neuron identity changed")
    require(graph["ptr"].shape == (NEURONS + 1,) and graph["ptr"][0] == 0 and
            graph["ptr"][-1] == EDGES and bool(np.all(np.diff(graph["ptr"]) >= 0)), "Invalid canonical CSR")
    for name in ("edge_id", "src", "dst", "signed_raw"):
        require(graph[name].shape == (EDGES,), "Graph edge count differs: " + name)
    require(np.array_equal(graph["src"], np.repeat(graph["nodes"], np.diff(graph["ptr"]))),
            "Canonical source rows disagree with CSR")
    for start in range(0, EDGES, 1 << 20):
        stop = min(start + (1 << 20), EDGES)
        require(np.array_equal(graph["edge_id"][start:stop], np.arange(start, stop, dtype=np.int64)),
                "Graph edge IDs are not canonical CSR row positions")
        require(bool(np.all((graph["dst"][start:stop] >= 0) & (graph["dst"][start:stop] < NEURONS))),
                "Target index outside complete graph")
    return graph


def load_inputs(case_dir, oracle_dir):
    with (case_dir / "spec.json").open(encoding="utf-8") as source:
        spec = json.load(source)
    for name in ("name", "steps", "dt_ms", "duration_ms", "seed_nodes", "model", "graph_dir"):
        require(name in spec, "Missing case field: " + name)
    steps = spec["steps"]
    seeds = spec["seed_nodes"]
    require(type(steps) is int and isinstance(seeds, list) and
            all(type(node) is int for node in seeds) and len(seeds) == len(set(seeds)), "Invalid case steps/seeds")
    seeds = tuple(sorted(seeds))
    require((steps, seeds) in ALLOWED_CASES, "Case not one of the three fixed full-graph trials")
    with (oracle_dir / "summary.json").open(encoding="utf-8") as source:
        oracle_summary = json.load(source)
    require(oracle_summary.get("complete") is True and oracle_summary.get("steps") == steps and
            oracle_summary.get("completed_ticks") == steps and oracle_summary.get("precision") == "float64" and
            oracle_summary.get("nodes") == NEURONS and oracle_summary.get("edges") == EDGES,
            "CPU oracle did not seal a complete trial with the declared steps/nodes/edges")
    require(float(spec["dt_ms"]) == .1 and float(spec["duration_ms"]) == steps * .1,
            "Unexpected duration or timestep")
    model = spec["model"]
    for name in ("resting_potential_mv", "threshold_mv", "reset_potential_mv", "synaptic_delay_ms",
                 "refractory_ms", "w_syn_mv_per_anatomical_synapse"):
        require(name in model and math.isfinite(float(model[name])), "Missing/nonfinite model field: " + name)
    require(float(model["synaptic_delay_ms"]) == 1.8 and float(model["refractory_ms"]) == 2.2 and
            float(model["w_syn_mv_per_anatomical_synapse"]) == .275, "Declared delay/refractory/W_syn changed")
    require(Path(spec["graph_dir"]).is_absolute(), "Case graph_dir must be absolute")
    graph = load_graph(str(Path(spec["graph_dir"]).resolve()))
    with (oracle_dir / "coefficients.json").open(encoding="utf-8") as source:
        coefficients = json.load(source)
    for name in COEFFICIENT_NAMES:
        require(name in coefficients and math.isfinite(float(coefficients[name])), "Missing/nonfinite oracle coefficient: " + name)
    require(coefficients["threshold"] == float(model["threshold_mv"]) * .001 and
            coefficients["reset"] == float(model["reset_potential_mv"]) * .001,
            "Oracle threshold/reset disagree with model declaration")
    histories = {}
    for name, dtype in (("v", np.float64), ("g", np.float64), ("ref", np.int32)):
        filename = name + "-history." + np.dtype(dtype).name + ".npy"
        value = np.load(oracle_dir / filename, mmap_mode="r", allow_pickle=False)
        require(value.shape == (steps, NEURONS) and value.dtype == dtype, "Invalid CPU state history: " + name)
        histories[name] = value
    weights = np.load(oracle_dir / "weights.npy", mmap_mode="r", allow_pickle=False)
    require(weights.dtype == np.float64 and weights.shape == (EDGES,) and bool(np.isfinite(weights).all()),
            "Invalid actual Brian2 SI weights")
    # A sign or row mismatch must not be hidden by copying the oracle weight file.
    for start in range(0, EDGES, 1 << 20):
        stop = min(start + (1 << 20), EDGES)
        require(np.array_equal(np.sign(weights[start:stop]), np.sign(graph["signed_raw"][start:stop])),
                "Oracle weight sign/zero pattern differs from graph")
    with np.load(oracle_dir / "state.npz", allow_pickle=False) as source:
        final = {name: source[name].copy() for name in source.files}
    for name, dtype in (("v", np.float64), ("g", np.float64), ("ref", np.int32), ("counts", np.int32)):
        require(final[name].dtype == dtype and final[name].shape == (NEURONS,), "Invalid CPU final state: " + name)
    require(final["step"].dtype == np.int32 and final["neuron_idx"].dtype == np.int32 and
            final["step"].ndim == 1 and final["step"].shape == final["neuron_idx"].shape, "Invalid CPU sparse events")
    require(bool(np.all((final["step"] >= 0) & (final["step"] < steps))) and
            bool(np.all(np.diff(final["step"]) >= 0)) and
            bool(np.all((final["neuron_idx"] >= 0) & (final["neuron_idx"] < NEURONS))), "CPU event index out of range")
    with (oracle_dir / "queue.jsonl").open(encoding="utf-8") as source:
        queues = [json.loads(line) for line in source if line.strip()]
    require(len(queues) == steps and [row["tick"] for row in queues] == list(range(steps)),
            "Missing, duplicate or reordered CPU queue records")
    with (oracle_dir / "requests.csv").open(encoding="utf-8", newline="") as source:
        reader = csv.DictReader(source)
        require(tuple(reader.fieldnames or ()) == REQUEST_FIELDS, "CPU request CSV schema mismatch")
        request_rows = [{name: int(row[name]) for name in REQUEST_FIELDS} for row in reader]
    require(len(request_rows) == len(seeds) and
            [row["neuron_idx"] for row in request_rows] == list(seeds) and
            all(row["tick"] == 0 and row["attempted"] == 1 and row["accepted"] in (0, 1) and
                row["rejected"] == 1 - row["accepted"] for row in request_rows), "CPU request coverage mismatch")
    return spec, graph, histories, weights, final, queues, coefficients, seeds, request_rows


def run_gpu(case_dir, out_dir, oracle_dir, build_root):
    """Execute one fixed full graph and compare every state bit, event and queue.

    Full CPU histories remain read-only. The first failing tick saves complete
    CPU/GPU V/g/ref arrays; every tick saves exact comparison fields and hashes.
    """
    case_dir, out_dir, oracle_dir = Path(case_dir), Path(out_dir), Path(oracle_dir)
    require(not out_dir.exists(), "Refusing existing output directory: " + str(out_dir))
    spec, graph, histories, weights, oracle_final, oracle_queues, coeff, seeds, oracle_requests = load_inputs(case_dir, oracle_dir)
    require(torch.cuda.is_available(), "CUDA unavailable")
    require(torch.cuda.device_count() == 1 and "4090" in torch.cuda.get_device_name(),
            "Only the declared single visible RTX4090 is permitted")
    out_dir.mkdir(parents=True, exist_ok=False)
    start = time.perf_counter()
    tick = -1
    rows = []
    first_failure = None
    try:
        backend = sealed_backend()
        ext, build_provenance = backend.build_extension(build_root)
        build_end = time.perf_counter()
        device = torch.device("cuda", torch.cuda.current_device())
        torch.cuda.reset_peak_memory_stats(device)

        def upload(value, dtype):
            # Writable copies prevent any accidental write to read-only source maps.
            return torch.from_numpy(np.array(value, copy=True, order="C")).to(device=device, dtype=dtype).contiguous()

        indptr = upload(graph["ptr"], torch.int64)
        dst = upload(graph["dst"], torch.int32)
        gpu_weights = upload(weights, torch.float64)
        rest = float(spec["model"]["resting_potential_mv"]) * .001
        v = torch.full((NEURONS,), rest, device=device, dtype=torch.float64)
        g = torch.zeros(NEURONS, device=device, dtype=torch.float64)
        ref = torch.zeros(NEURONS, device=device, dtype=torch.int32)
        counts = torch.zeros(NEURONS, device=device, dtype=torch.int32)
        requests = torch.zeros(NEURONS, device=device, dtype=torch.uint8)
        if seeds:
            requests[torch.tensor(seeds, device=device, dtype=torch.int64)] = 1
        ring = torch.zeros((19, NEURONS), device=device, dtype=torch.uint8)
        events_step, events_node = [], []
        requested_total = accepted_request_total = 0
        total_due = total_accepted = 0
        torch.cuda.synchronize(device)
        simulation_start = time.perf_counter()
        with (out_dir / "step-comparisons.jsonl").open("x", encoding="utf-8", buffering=1) as comparisons, \
             (out_dir / "queue.jsonl").open("x", encoding="utf-8", buffering=1) as queue_output, \
             (out_dir / "tick-state-hashes.jsonl").open("x", encoding="utf-8", buffering=1) as hash_output, \
             (out_dir / "requests.csv").open("x", encoding="utf-8", newline="", buffering=1) as request_output:
            request_writer = csv.DictWriter(request_output, fieldnames=REQUEST_FIELDS)
            request_writer.writeheader()
            for tick in range(spec["steps"]):
                if tick == 1:
                    requests.zero_()
                requested = int(torch.count_nonzero(requests).item())
                accepted_requests = int(torch.count_nonzero((requests != 0) & (ref <= 1)).item())
                requested_total += requested
                accepted_request_total += accepted_requests
                request_rows = []
                if requested:
                    requested_nodes = torch.nonzero(requests, as_tuple=False).flatten()
                    accepted_mask = (ref[requested_nodes] <= 1).to(torch.uint8).cpu().numpy()
                    for node, accepted in zip(requested_nodes.cpu().numpy(), accepted_mask):
                        request_rows.append({"tick": tick, "neuron_idx": int(node), "attempted": 1,
                                             "accepted": int(accepted), "rejected": 1 - int(accepted)})
                    request_writer.writerows(request_rows)
                requests_exact = request_rows == [row for row in oracle_requests if row["tick"] == tick]
                spikes = ring[tick % 19]
                ext.state_step(v, g, ref, requests, spikes, counts,
                               *[float(coeff[name]) for name in COEFFICIENT_NAMES], 22)
                expanded_cpu = np.empty(0, dtype=np.int64)
                accepted_cpu = np.empty(0, dtype=np.uint8)
                if tick >= 18:
                    due = ring[(tick - 18) % 19]
                    sources = torch.nonzero(due, as_tuple=False).flatten().contiguous()
                    if sources.numel():
                        degrees = indptr[sources + 1] - indptr[sources]
                        offsets = (torch.cumsum(degrees, dim=0) - degrees).contiguous()
                        edge_count = int(degrees.sum().item())
                        expanded = torch.empty(edge_count, device=device, dtype=torch.int64)
                        ext.fill_edges(indptr, sources, offsets, expanded)
                        if edge_count:
                            destination = dst[expanded]
                            # Hash the actual GPU-expanded queue BEFORE target grouping.
                            expanded_cpu = expanded.cpu().numpy().copy()
                            accepted_cpu = (ref[destination.long()] == 0).to(torch.uint8).cpu().numpy().copy()
                            order = torch.argsort(destination, stable=True)
                            ordered_edges = expanded[order].contiguous()
                            ordered_targets = destination[order]
                            boundary = torch.ones(edge_count, device=device, dtype=torch.bool)
                            boundary[1:] = ordered_targets[1:] != ordered_targets[:-1]
                            starts = torch.nonzero(boundary, as_tuple=False).flatten().contiguous()
                            targets = ordered_targets[starts].contiguous()
                            ext.apply_arrivals(ordered_edges, starts, targets, gpu_weights, g, ref)
                queue = queue_summary(tick, expanded_cpu, accepted_cpu)
                queue_output.write(json.dumps(queue, allow_nan=False) + "\n")
                total_due += queue["count"]
                total_accepted += queue["accepted_count"]
                queue_exact = queue == oracle_queues[tick]
                # No external input exists in these fixed cases after the tick-zero requests.
                actual = {"v": v.cpu().numpy().copy(), "g": g.cpu().numpy().copy(),
                          "ref": ref.cpu().numpy().copy()}
                state_checks = {name: bitwise_comparison(actual[name], histories[name][tick], name != "ref")
                                for name in ("v", "g", "ref")}
                fired = torch.nonzero(spikes, as_tuple=False).flatten().cpu().numpy().astype(np.int32)
                if len(fired):
                    events_step.append(np.full(len(fired), tick, dtype=np.int32))
                    events_node.append(fired.copy())
                oracle_nodes = oracle_final["neuron_idx"][oracle_final["step"] == tick]
                events_exact = bool(np.array_equal(fired, oracle_nodes))
                finite = state_checks["v"]["finite"] and state_checks["g"]["finite"]
                passed = bool(finite and queue_exact and events_exact and requests_exact and
                              all(value["bitwise_exact"] for value in state_checks.values()))
                row = {"tick": tick, "pass_all": passed, "states": state_checks,
                       "events_exact": events_exact, "gpu_event_count": len(fired), "cpu_event_count": len(oracle_nodes),
                       "queue_exact": queue_exact, "gpu_queue": queue, "cpu_queue": oracle_queues[tick],
                       "requests": requested, "accepted_requests": accepted_requests,
                       "rejected_requests": requested - accepted_requests,
                       "request_rows_exact": requests_exact,
                       "state_sha256": {name: array_digest(actual[name], "<i4" if name == "ref" else "<f8")
                                        for name in ("v", "g", "ref")}}
                comparisons.write(json.dumps(row, allow_nan=False) + "\n")
                hash_output.write(json.dumps({"tick": tick,
                    **{name + "_sha256": value for name, value in row["state_sha256"].items()}},
                    allow_nan=False) + "\n")
                rows.append(row)
                if not passed and first_failure is None:
                    first_failure = row
                    with (out_dir / "first-failure-state.npz").open("xb") as target:
                        np.savez_compressed(target, tick=np.array(tick, dtype=np.int32),
                            **{"gpu_" + name: value for name, value in actual.items()},
                            **{"cpu_" + name: np.asarray(histories[name][tick]).copy() for name in actual},
                            gpu_event_nodes=fired, cpu_event_nodes=oracle_nodes,
                            arrival_edge_id=expanded_cpu, arrival_accepted=accepted_cpu)
                    write_json(out_dir / "first-failure.json", row)
        torch.cuda.synchronize(device)
        simulation_end = time.perf_counter()

        def join(values):
            return np.concatenate(values).astype(np.int32, copy=False) if values else np.empty(0, dtype=np.int32)

        state = {"step": join(events_step), "neuron_idx": join(events_node), "counts": counts.cpu().numpy(),
                 "v": v.cpu().numpy(), "g": g.cpu().numpy(), "ref": ref.cpu().numpy()}
        with (out_dir / "state.npz").open("xb") as target:
            np.savez_compressed(target, **state)
        final_checks = {name: bool(np.array_equal(value.view(np.uint64), oracle_final[name].view(np.uint64)))
                        if name in ("v", "g") else bool(np.array_equal(value, oracle_final[name]))
                        for name, value in state.items()}
        history_final_checks = {name: bool(np.array_equal(oracle_final[name].view(np.uint64), histories[name][-1].view(np.uint64)))
                                if name in ("v", "g") else bool(np.array_equal(oracle_final[name], histories[name][-1]))
                                for name in ("v", "g", "ref")}
        pass_all = bool(len(rows) == spec["steps"] and all(row["pass_all"] for row in rows) and
                        all(final_checks.values()) and all(history_final_checks.values()) and
                        requested_total == len(seeds) and accepted_request_total == len(seeds))
        summary = {"case": spec["name"], "pass_all": pass_all, "precision": "float64", "state_units": "SI volts",
                   "neurons": NEURONS, "edges": EDGES, "steps_expected": spec["steps"], "steps_compared": len(rows),
                   "all_neurons_compared_every_tick": True, "bitwise_includes_signed_zero": True,
                   "first_failed_tick": None if first_failure is None else first_failure["tick"],
                   "failed_tick_count": sum(not row["pass_all"] for row in rows),
                   "final_state_exact": final_checks, "oracle_final_matches_history": history_final_checks,
                   "total_spikes": len(state["step"]), "arriving_edges": total_due, "accepted_arriving_edges": total_accepted,
                   "requested": requested_total, "accepted_requests": accepted_request_total,
                   "rejected_requests": requested_total - accepted_request_total,
                   "whole_graph_numerical_comparison": True, "olfactory_representation_validated": False,
                   "apl_overlay": False, "learning": False, "candidate_generated": False,
                   "performance_runtime_claim": False, "cpu_sync_each_tick": True,
                   "first_failure_full_state_saved": first_failure is not None,
                   "full_gpu_history_saved": False, "sealed_4AE_model": str(AE_MODEL),
                   "build": build_provenance, "build_or_cache_seconds": build_end - start,
                   "diagnostic_simulation_seconds": simulation_end - simulation_start,
                   "wall_seconds": time.perf_counter() - start, "gpu": torch.cuda.get_device_name(device),
                   "peak_allocated_vram_bytes": int(torch.cuda.max_memory_allocated(device)),
                   "peak_reserved_vram_bytes": int(torch.cuda.max_memory_reserved(device)),
                   "input_sha256": {"spec.json": digest(case_dir / "spec.json"),
                                    "oracle/weights.npy": digest(oracle_dir / "weights.npy"),
                                    "oracle/coefficients.json": digest(oracle_dir / "coefficients.json")}}
        write_json(out_dir / "summary.json", summary)
        write_json(out_dir / "metrics.json", {name: summary[name] for name in
                   ("wall_seconds", "build_or_cache_seconds", "diagnostic_simulation_seconds", "gpu",
                    "peak_allocated_vram_bytes", "peak_reserved_vram_bytes", "performance_runtime_claim", "cpu_sync_each_tick")})
        print(json.dumps({"case": spec["name"], "pass_all": pass_all, "steps": len(rows),
                          "first_failed_tick": summary["first_failed_tick"], "total_spikes": len(state["step"])}), flush=True)
        return summary
    except Exception as error:
        write_json(out_dir / "error.json", {"type": type(error).__name__, "message": str(error),
                   "traceback": traceback.format_exc(), "tick_in_progress": tick,
                   "completed_tick_comparisons": len(rows), "pass_all": False})
        raise