"""Diagnostic author-float64 LIF plus frozen legacy-float32 APL overlay.

The sealed Phase 4AE CUDA operators are reused without modification.  APL is
an explicitly mixed-precision engineering overlay, not an author-model claim.
All CUDA/CPU synchronization here is diagnostic; no performance claim is made.
"""
from functools import lru_cache
import csv
import hashlib
import json
import math
from pathlib import Path
import time
import traceback
import types

import numpy as np
import torch

from legacy_apl_worker import LegacyAPL, N_NEURONS, file_sha256, array_sha256


SOURCE_ROOT = Path(__file__).resolve().parent
AE_MODEL = SOURCE_ROOT.parent / "author-runtime-lab-v1"
AF_MODEL = SOURCE_ROOT.parent / "author-fullgraph-lab-v1"
COEFFICIENT_NAMES = ("gd", "c0", "cg", "cv", "threshold", "reset", "seed_volt")
REQUEST_FIELDS = ("tick", "neuron_idx", "attempted", "accepted", "rejected")


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 raw_hash(value, dtype):
    return hashlib.sha256(np.ascontiguousarray(value, dtype=dtype).tobytes()).hexdigest()


def load_source(path, name):
    module = types.ModuleType(name)
    module.__file__ = str(path)
    exec(compile(Path(path).read_bytes(), str(path), "exec"), module.__dict__)
    return module


@lru_cache(maxsize=1)
def sealed_backend():
    return load_source(AE_MODEL / "gpu_backend.py", "phase4ag_sealed_ae_backend")


@lru_cache(maxsize=1)
def queue_function():
    return load_source(AF_MODEL / "oracle_full.py", "phase4ag_sealed_af_queue").queue_summary


@lru_cache(maxsize=4)
def load_graph(graph_directory):
    root = Path(graph_directory)
    require(root.is_absolute(), "graph_dir must be absolute")
    graph = {}
    for name, dtype in (("nodes", np.int32), ("edge_id", np.int64), ("src", np.int32),
                        ("dst", np.int32), ("signed_raw", np.int64), ("ptr", np.int64),
                        ("weights", np.float64)):
        value = np.load(root / (name + ".npy"), mmap_mode="r", allow_pickle=False)
        require(value.dtype == dtype and value.ndim == 1, "Graph dtype/rank mismatch: " + name)
        graph[name] = value
    nodes, edges = graph["nodes"], graph["edge_id"]
    n, e = len(nodes), len(edges)
    require(0 < n <= N_NEURONS and bool(np.all((nodes >= 0) & (nodes < N_NEURONS))) and
            bool(np.all(np.diff(nodes) > 0)), "Global node IDs must be sorted unique")
    require(n < 512 or (n == N_NEURONS and np.array_equal(nodes, np.arange(N_NEURONS, dtype=np.int32))),
            "Only small fixtures or the intact frozen full graph are permitted")
    require(n != N_NEURONS or e == 25582938, "Full graph edge count changed")
    for name in ("src", "dst", "signed_raw", "weights"):
        require(graph[name].shape == (e,), "Graph edge array length mismatch: " + name)
    ptr = graph["ptr"]
    require(ptr.shape == (n + 1,) and ptr[0] == 0 and ptr[-1] == e and
            bool(np.all(np.diff(ptr) >= 0)), "Invalid graph CSR")
    require(np.array_equal(graph["src"], np.repeat(np.arange(n, dtype=np.int32), np.diff(ptr))),
            "Source order does not match CSR")
    # Global edge IDs need not equal local row indices in a fixture.
    for start in range(0, e, 1 << 20):
        stop = min(start + (1 << 20), e)
        require(bool(np.all((graph["dst"][start:stop] >= 0) & (graph["dst"][start:stop] < n))),
                "Local destination out of range")
        require(bool(np.isfinite(graph["weights"][start:stop]).all()), "Nonfinite base SI weights")
        require(np.array_equal(np.sign(graph["weights"][start:stop]), np.sign(graph["signed_raw"][start:stop])),
                "Base weight sign/zero pattern differs from frozen raw contacts")
        lo = max(0, start - 1)
        require(bool(np.all(edges[start:stop] >= 0)) and bool(np.all(np.diff(edges[lo:stop]) > 0)),
                "Global edge IDs must retain canonical increasing source/edge order")
    return graph


def load_inputs(case_dir):
    spec = json.loads((case_dir / "spec.json").read_text(encoding="utf-8"))
    for name in ("name", "steps", "dt_ms", "duration_ms", "model", "graph_dir", "apl",
                 "apl_bundle_dir", "coefficients_path", "record_history"):
        require(name in spec, "Missing case field: " + name)
    steps = spec["steps"]
    require(type(steps) is int and 0 < steps <= 1000 and float(spec["dt_ms"]) == .1 and
            float(spec["duration_ms"]) == steps * .1, "Unexpected steps/dt/duration")
    require(type(spec["record_history"]) is bool, "record_history must be Boolean")
    require(spec["apl"] == {"tau_ms": 20, "gain": .01, "normalization": 1}, "Frozen APL parameters changed")
    for name in ("graph_dir", "apl_bundle_dir", "coefficients_path"):
        require(Path(spec[name]).is_absolute(), "Input path must be absolute: " + name)
    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])), "Invalid 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, "Frozen delay/ref/W_syn changed")
    coeff = json.loads(Path(spec["coefficients_path"]).read_text(encoding="utf-8"))
    for name in COEFFICIENT_NAMES:
        require(name in coeff and math.isfinite(float(coeff[name])), "Invalid author coefficient: " + name)
    require(coeff["threshold"] == float(model["threshold_mv"]) * .001 and
            coeff["reset"] == float(model["reset_potential_mv"]) * .001, "Coefficient threshold/reset mismatch")
    graph = load_graph(str(Path(spec["graph_dir"]).resolve()))
    n = len(graph["nodes"])
    require(spec["record_history"] == (n < 512), "Fixtures need dense history; full graph must not duplicate dense histories")
    with np.load(case_dir / "requests.npz", allow_pickle=False) as source:
        require(set(source.files) == {"tick", "neuron_idx"}, "Request sparse schema changed")
        ticks, requested = source["tick"].copy(), source["neuron_idx"].copy()
    require(ticks.dtype == np.int32 and requested.dtype == np.int32 and ticks.ndim == 1 and
            ticks.shape == requested.shape, "Invalid sparse request dtype/shape")
    require(bool(np.all((ticks >= 0) & (ticks < steps))), "Request tick outside trial")
    if len(ticks) > 1:
        require(bool(np.all((ticks[1:] > ticks[:-1]) |
                           ((ticks[1:] == ticks[:-1]) & (requested[1:] > requested[:-1])))),
                "Requests must be sorted unique by tick/global neuron")
    local = np.searchsorted(graph["nodes"], requested)
    require(bool(np.all(local < n)) and np.array_equal(graph["nodes"][local], requested),
            "Requested neuron absent from graph")
    initial = {"initial_v_mv": np.full(n, float(model["resting_potential_mv"]), dtype=np.float64),
               "initial_g_mv": np.zeros(n, dtype=np.float64)}
    initial_path = case_dir / "arrays.npz"
    if initial_path.exists():
        with np.load(initial_path, allow_pickle=False) as source:
            require(set(source.files).issubset(initial), "Unexpected initial-state array")
            for name in source.files:
                value = source[name]
                require(value.dtype == np.float64 and value.shape == (n,) and bool(np.isfinite(value).all()),
                        "Invalid initial state: " + name)
                initial[name] = value.copy()
    return spec, graph, coeff, ticks, requested, local, initial


def new_array(path, dtype, shape):
    require(not Path(path).exists(), "Refusing existing output: " + str(path))
    return np.lib.format.open_memmap(path, mode="w+", dtype=dtype, shape=shape)


def run_gpu(case_dir, out_dir, build_root):
    """Run fresh candidate state; never consume oracle/candidate APL traces.

    pass_all denotes internal recording/finite integrity only.  Independent
    comparison and the original complete representation gate are separate.
    """
    case_dir, out_dir = Path(case_dir), Path(out_dir)
    require(not out_dir.exists(), "Refusing existing output directory: " + str(out_dir))
    out_dir.mkdir(parents=True, exist_ok=False)
    start, tick, completed = time.perf_counter(), -1, 0
    maps = {}
    try:
        spec, graph, coeff, request_ticks, request_nodes, request_local, initial = load_inputs(case_dir)
        require(torch.cuda.is_available() and torch.cuda.device_count() == 1 and
                "4090" in torch.cuda.get_device_name(), "Only the declared single visible RTX4090 is permitted")
        ext, build = sealed_backend().build_extension(build_root)
        build_done = time.perf_counter()
        apl = LegacyAPL(spec["apl_bundle_dir"], spec["dt_ms"])
        write_json(out_dir / "apl-component.json", apl.metadata)
        device = apl.device
        torch.cuda.reset_peak_memory_stats(device)
        n, e, steps = len(graph["nodes"]), len(graph["edge_id"]), spec["steps"]

        def upload(value, dtype):
            return torch.from_numpy(np.array(value, copy=True, order="C")).to(device=device, dtype=dtype).contiguous()

        indptr, dst = upload(graph["ptr"], torch.int64), upload(graph["dst"], torch.int32)
        global_nodes_gpu = upload(graph["nodes"], torch.int64)
        base_weights = upload(graph["weights"], torch.float64)
        weights = base_weights.clone()
        del base_weights
        local_apl = np.flatnonzero(np.isin(graph["nodes"], apl.apl_globals)).astype(np.int64)
        apl_out_mask = np.isin(graph["src"], local_apl)
        if bool(np.any(apl_out_mask)):
            weights[upload(apl_out_mask, torch.bool)] = 0.0
        effective_hash = hashlib.sha256()
        for offset in range(0, e, 1 << 20):
            effective_hash.update(weights[offset:offset + (1 << 20)].cpu().numpy().astype("<f8", copy=False).tobytes())
        common_kc, local_kc, full_kc = np.intersect1d(graph["nodes"], apl.kc_globals, assume_unique=True, return_indices=True)
        kc_local_gpu, kc_full_gpu = upload(local_kc, torch.int64), upload(full_kc, torch.int64)
        v = upload(initial["initial_v_mv"] * .001, torch.float64)
        g = upload(initial["initial_g_mv"] * .001, torch.float64)
        ref = torch.zeros(n, device=device, dtype=torch.int32)
        counts = torch.zeros(n, device=device, dtype=torch.int32)
        requests = torch.zeros(n, device=device, dtype=torch.uint8)
        ring = torch.zeros((19, n), device=device, dtype=torch.uint8)
        for name, width in (("drive", 2), ("state", 2), ("inhibition", 4064)):
            maps["apl_" + name] = new_array(out_dir / ("apl-" + name + ".float32.npy"), np.float32, (steps, width))
        if spec["record_history"]:
            for name, dtype in (("v", np.float64), ("g", np.float64), ("ref", np.int32)):
                maps[name] = new_array(out_dir / (name + "-history." + np.dtype(dtype).name + ".npy"), dtype, (steps, n))
                maps["pre_" + name] = new_array(out_dir / ("preAPL-" + name + "." + np.dtype(dtype).name + ".npy"), dtype, (steps, n))
        event_steps, event_nodes = [], []
        accepted_requests_total, due_total, accepted_due_total = 0, 0, 0
        queue_summary = queue_function()
        simulation_start = time.perf_counter()
        with (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:
            writer = csv.DictWriter(request_output, fieldnames=REQUEST_FIELDS)
            writer.writeheader()
            for tick in range(steps):
                requests.zero_()
                begin, end = np.searchsorted(request_ticks, (tick, tick + 1))
                if end > begin:
                    requested_local_gpu = upload(request_local[begin:end], torch.int64)
                    requests[requested_local_gpu] = 1
                    accepted = (ref[requested_local_gpu] <= 1).to(torch.uint8).cpu().numpy()
                    accepted_requests_total += int(accepted.sum())
                    writer.writerows({"tick": tick, "neuron_idx": int(node), "attempted": 1,
                                      "accepted": int(ok), "rejected": 1 - int(ok)}
                                     for node, ok in zip(request_nodes[begin:end], accepted))
                spikes = ring[tick % 19]
                ext.state_step(v, g, ref, requests, spikes, counts,
                               *[float(coeff[name]) for name in COEFFICIENT_NAMES], 22)
                expanded_global = 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]
                            # Actual active fanout order, BEFORE stable target grouping.
                            expanded_global = graph["edge_id"][expanded.cpu().numpy()].copy()
                            accepted_cpu = (ref[destination.long()] == 0).to(torch.uint8).cpu().numpy().copy()
                            order = torch.argsort(destination, stable=True)
                            ordered_edges, ordered_targets = expanded[order].contiguous(), 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, weights, g, ref)
                queue = queue_summary(tick, expanded_global, accepted_cpu)
                queue_output.write(json.dumps(queue, allow_nan=False) + "\n")
                due_total += queue["count"]
                accepted_due_total += queue["accepted_count"]
                if spec["record_history"]:
                    for name, value in (("v", v), ("g", g), ("ref", ref)):
                        maps["pre_" + name][tick] = value.cpu().numpy()
                if n == N_NEURONS:
                    full_spikes = spikes
                else:
                    apl.spikes.zero_()
                    apl.spikes[global_nodes_gpu] = spikes
                    full_spikes = apl.spikes
                drive, apl_state, h = apl.advance(full_spikes)
                # Mixed precision contract: widen f32 h EXACTLY, THEN multiply
                # in f64, THEN subtract in f64.  Do not guard refractory KCs.
                delta_si = h[kc_full_gpu].to(torch.float64) * 0.001
                g[kc_local_gpu] = g[kc_local_gpu] - delta_si
                apl_cpu = {"drive": drive.cpu().numpy().copy(), "state": apl_state.cpu().numpy().copy(),
                           "inhibition": h.cpu().numpy().copy()}
                for name, value in apl_cpu.items():
                    maps["apl_" + name][tick] = value
                actual = {"v": v.cpu().numpy().copy(), "g": g.cpu().numpy().copy(),
                          "ref": ref.cpu().numpy().copy()}
                finite = all(bool(np.isfinite(value).all()) for value in list(actual.values()) + list(apl_cpu.values()))
                row = {"tick": tick, **{name + "_sha256": raw_hash(value, "<i4" if name == "ref" else "<f8")
                                         for name, value in actual.items()}}
                hash_output.write(json.dumps(row, allow_nan=False) + "\n")
                if spec["record_history"]:
                    for name, value in actual.items():
                        maps[name][tick] = value
                fired_local = torch.nonzero(spikes, as_tuple=False).flatten().cpu().numpy()
                if len(fired_local):
                    event_steps.append(np.full(len(fired_local), tick, dtype=np.int32))
                    event_nodes.append(graph["nodes"][fired_local].copy())
                completed += 1
                if not finite:
                    with (out_dir / "first-failure-state.npz").open("xb") as target:
                        np.savez_compressed(target, tick=np.array(tick, dtype=np.int32), **actual,
                                            **{"apl_" + name: value for name, value in apl_cpu.items()})
                    require(False, "Nonfinite state; first failure saved")
        torch.cuda.synchronize(device)

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

        state = {"step": joined(event_steps), "neuron_idx": joined(event_nodes), "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)
        event_locals = np.searchsorted(graph["nodes"], state["neuron_idx"])
        counts_exact = np.array_equal(np.bincount(event_locals, minlength=n), state["counts"])
        require(counts_exact and completed == steps, "Sparse events/counts/completed ticks mismatch")
        for value in maps.values():
            value.flush()
        summary = {"case": spec["name"], "complete": True, "pass_all": True,
                   "pass_all_scope": "internal finite/recording/event-count integrity only; independent equivalence not claimed",
                   "steps": steps, "completed_ticks": completed, "nodes": n, "edges": e, "precision": "float64",
                   "apl_precision": "float32", "state_units": "SI volts", "total_spikes": len(state["step"]),
                   "requested": len(request_nodes), "accepted_requests": accepted_requests_total,
                   "rejected_requests": len(request_nodes) - accepted_requests_total,
                   "arriving_edges": due_total, "accepted_arriving_edges": accepted_due_total,
                   "counts_match_sparse_events": bool(counts_exact), "all_finite": True,
                   "record_history": spec["record_history"], "apl_fast_edges_zeroed": int(apl_out_mask.sum()),
                   "kc_targets_present": len(common_kc), "effective_weights_raw_sha256": effective_hash.hexdigest(),
                   "apl_component": apl.metadata, "apl_schedule": "end, after threshold/reset and guarded arrivals, unconditional all KC",
                   "apl_si_conversion": "h32 -> float64 -> multiply float64 0.001 -> separate float64 subtraction",
                   "build": build, "sealed_ae_backend_sha256": file_sha256(AE_MODEL / "gpu_backend.py"),
                   "shared_apl_component_oracle_limitation": True, "oracle_validation_claim": False,
                   "learning": False, "performance_runtime_claim": False, "cpu_sync_each_tick": True,
                   "build_or_cache_seconds": build_done - start,
                   "diagnostic_simulation_seconds": time.perf_counter() - simulation_start,
                   "wall_seconds": time.perf_counter() - start,
                   "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": file_sha256(case_dir / "spec.json"),
                                    "requests.npz": file_sha256(case_dir / "requests.npz"),
                                    "base_weights.npy": file_sha256(Path(spec["graph_dir"]) / "weights.npy"),
                                    "coefficients.json": file_sha256(spec["coefficients_path"])}}
        if (case_dir / "arrays.npz").exists():
            summary["input_sha256"]["arrays.npz"] = file_sha256(case_dir / "arrays.npz")
        write_json(out_dir / "summary.json", summary)
        return summary
    except Exception as error:
        for value in maps.values():
            value.flush()
        write_json(out_dir / "error.json", {"complete": False, "pass_all": False,
                   "tick_in_progress": tick, "completed_ticks": completed,
                   "type": type(error).__name__, "message": str(error), "traceback": traceback.format_exc()})
        raise