"""Independent CUDA eligibility candidate; never installs weights or runs a brain.

The caller authenticates frozen sources, configuration, the complete KC list,
and the candidate edge mapping. This module does not infer D from DAN activity.
CuPy is imported only when CudaEligibility is instantiated. Set CUPY_CACHE_DIR
to a dedicated candidate directory before doing so; this module writes no files.

API: prepare_events(...) once, then CudaEligibility(...).run(prepared, ...).
All commands are sampled AFTER the current tick; only the LAST command snapshot
feeds the single epoch-boundary update. With no commands sampled_trace is +0.
cpu_reference has the same run arguments for host-only validation, and is not a
replacement for comparison against the SHA-pinned frozen implementation.
"""
from __future__ import annotations

from dataclasses import dataclass
import math
import os
from pathlib import Path
import time
from typing import Any

import numpy as np


COMPARTMENTS = ("B2B'2a", "a1", "y2a'1", "y5B'2a")
_INT32_MAX = np.iinfo(np.int32).max
CUDA_OPTIONS = ("--std=c++11", "--fmad=false", "--ftz=false")
CUDA_SOURCE = r'''
extern "C" __global__
void eligibility_trace(const long long* ptr, const int* ticks,
                       const int* commands, const int n_commands,
                       const int n_kc, const int n_steps, const float decay,
                       float* snapshots, float* final_trace, float* sampled) {
    const int k = blockDim.x * blockIdx.x + threadIdx.x;
    if (k >= n_kc) return;
    long long cursor = ptr[k];
    const long long end = ptr[k + 1];
    int command = 0;
    float trace = 0.0f;
    float last_sample = 0.0f;
    // Each KC follows the FULL sequential recurrence, including empty ticks.
    for (int tick = 0; tick < n_steps; ++tick) {
        float spike = 0.0f;
        if (cursor < end && ticks[cursor] == tick) {
            spike = 1.0f;
            ++cursor;
        }
        trace = __fmul_rn(trace, decay);
        trace = __fadd_rn(trace, spike);
        trace = trace < 1.0f ? trace : 1.0f;
        if (command < n_commands && commands[command] == tick) {
            snapshots[(long long)command * n_kc + k] = trace;
            last_sample = trace;
            ++command;
        }
    }
    final_trace[k] = trace;
    sampled[k] = last_sample;
}

extern "C" __global__
void eligibility_delta(const float* sampled, const int* source_local,
                       const unsigned char* selected, const float* delta,
                       const float* base, const int n_edges,
                       const float negative_eta_times_d,
                       const float negative_floor, const int nonzero_d,
                       const int frozen, float* output) {
    const int edge = blockDim.x * blockIdx.x + threadIdx.x;
    if (edge >= n_edges) return;
    if (frozen) {
        output[edge] = delta[edge];
        return;
    }
    float update = 0.0f;
    if (nonzero_d && selected[edge])
        update = __fmul_rn(negative_eta_times_d, sampled[source_local[edge]]);
    const float out = __fadd_rn(delta[edge], update);
    const float lower = __fmul_rn(negative_floor, base[edge]);
    // np.maximum chooses its SECOND operand on equal finite values, including
    // +0/-0. CUDA fmaxf is deliberately not substituted for that tie behavior.
    output[edge] = out > lower ? out : lower;
}
'''


def _require(condition: Any, message: str) -> None:
    if not condition:
        raise ValueError(message)


def _integer(value: Any, name: str, minimum: int = 1) -> int:
    _require(isinstance(value, (int, np.integer)) and not isinstance(value, (bool, np.bool_)),
             f"{name} must be an integer")
    _require(minimum <= int(value) <= _INT32_MAX, f"{name} outside supported int32 range")
    return int(value)


def _array(value: Any, dtype: Any, name: str) -> np.ndarray:
    _require(isinstance(value, np.ndarray) and value.ndim == 1 and value.dtype == np.dtype(dtype),
             f"{name} must be a one-dimensional {np.dtype(dtype)} ndarray")
    return value


def _owned(value: np.ndarray) -> np.ndarray:
    result = np.array(value, copy=True, order="C")
    result.setflags(write=False)
    return result


def _scalar(value: Any, name: str) -> float:
    _require(isinstance(value, (float, int, np.floating, np.integer))
             and not isinstance(value, (bool, np.bool_)), f"{name} must be a real scalar")
    try:
        result = float(value)
    except (ValueError, OverflowError) as exc:
        raise ValueError(f"{name} must be a finite scalar") from exc
    _require(math.isfinite(result), f"{name} must be finite")
    return result


@dataclass(frozen=True)
class PreparedEvents:
    kc: np.ndarray
    indptr: np.ndarray
    ticks: np.ndarray
    command_steps: np.ndarray
    steps: int
    n_neurons: int
    input_event_count: int
    preparation_seconds: float

    @property
    def kc_event_count(self) -> int:
        return len(self.ticks)


@dataclass(frozen=True)
class EligibilityResult:
    sampled_trace: np.ndarray
    final_trace: np.ndarray
    snapshots: np.ndarray
    delta_after: np.ndarray
    timing: dict[str, Any]
    audit: dict[str, Any]


def prepare_events(step: np.ndarray, neuron_idx: np.ndarray, kc: np.ndarray,
                   steps: int, command_steps: Any, *, n_neurons: int = 166700) -> PreparedEvents:
    """Validate original sparse events and pack sorted per-KC event CSR on CPU.

    Inputs must be sorted by (tick, global neuron ID), unique and int32. No
    [steps, neurons] raster is allocated. Original input storage is not retained.
    Empty, one, or multiple sorted unique command ticks are supported.
    """
    started = time.perf_counter()
    steps = _integer(steps, "steps")
    n_neurons = _integer(n_neurons, "n_neurons")
    step = _array(step, np.int32, "step")
    neuron_idx = _array(neuron_idx, np.int32, "neuron_idx")
    kc = _array(kc, np.int32, "kc")
    _require(0 < len(kc) <= _INT32_MAX and np.all((kc >= 0) & (kc < n_neurons)),
             "KC IDs empty or outside graph")
    _require(np.all(kc[1:] > kc[:-1]), "KC IDs must be sorted unique global IDs")
    _require(len(step) == len(neuron_idx), "event columns have different lengths")
    _require(np.all((step >= 0) & (step < steps)), "event tick outside trial")
    _require(np.all((neuron_idx >= 0) & (neuron_idx < n_neurons)), "event neuron outside graph")
    _require(np.all((step[1:] > step[:-1]) |
                    ((step[1:] == step[:-1]) & (neuron_idx[1:] > neuron_idx[:-1]))),
             "events must be sorted by tick/global ID and unique")
    commands = np.asarray(command_steps)
    _require(commands.ndim == 1, "command_steps must be one-dimensional")
    if commands.size == 0:
        commands = np.empty(0, np.int32)
    _require(commands.dtype.kind in "iu", "command_steps must be integers")
    _require(np.all((commands >= 0) & (commands < steps)), "command tick outside trial")
    _require(np.all(commands[1:] > commands[:-1]), "command_steps must be sorted unique")
    local = np.searchsorted(kc, neuron_idx)
    keep = (local < len(kc)) & (kc[np.minimum(local, len(kc) - 1)] == neuron_idx)
    local = local[keep]
    order = np.argsort(local, kind="stable")
    ticks = step[keep][order]
    counts = np.bincount(local, minlength=len(kc)).astype(np.int64, copy=False)
    ptr = np.empty(len(kc) + 1, np.int64)
    ptr[0] = 0
    np.cumsum(counts, out=ptr[1:])
    return PreparedEvents(_owned(kc), _owned(ptr), _owned(ticks),
                          _owned(commands.astype(np.int32)), steps, n_neurons,
                          len(step), time.perf_counter() - started)


def _validate_prepared(prepared: PreparedEvents) -> None:
    _require(isinstance(prepared, PreparedEvents), "prepared must come from prepare_events")
    _integer(prepared.steps, "steps")
    _integer(prepared.n_neurons, "n_neurons")
    kc = _array(prepared.kc, np.int32, "prepared.kc")
    ptr = _array(prepared.indptr, np.int64, "prepared.indptr")
    ticks = _array(prepared.ticks, np.int32, "prepared.ticks")
    commands = _array(prepared.command_steps, np.int32, "prepared.command_steps")
    _require(0 < len(kc) <= _INT32_MAX and np.all((kc >= 0) & (kc < prepared.n_neurons))
             and np.all(kc[1:] > kc[:-1]), "invalid prepared KC IDs")
    _require(len(ptr) == len(kc) + 1 and ptr[0] == 0 and ptr[-1] == len(ticks)
             and np.all(ptr[1:] >= ptr[:-1]), "invalid prepared CSR offsets")
    _require(np.all((ticks >= 0) & (ticks < prepared.steps)), "invalid prepared event ticks")
    # A decrease/equality is permitted only at a CSR row boundary. Checking all
    # adjacent pairs stays vectorized and also catches a forged PreparedEvents.
    nonincreasing = ticks[1:] <= ticks[:-1]
    if len(ticks) > 1:
        boundaries = ptr[1:-1]
        boundaries = boundaries[(boundaries > 0) & (boundaries < len(ticks))] - 1
        nonincreasing[boundaries] = False
    _require(not np.any(nonincreasing), "prepared KC ticks must be sorted unique")
    _require(np.all((commands >= 0) & (commands < prepared.steps))
             and np.all(commands[1:] > commands[:-1]), "invalid prepared command ticks")


def _validate_run(prepared, candidate_sources, candidate_compartments, compartment,
                  delta_before, base_candidate, d, frozen, eta_mv, floor_fraction, decay):
    _validate_prepared(prepared)
    sources = _array(candidate_sources, np.int32, "candidate_sources")
    delta = _array(delta_before, np.float32, "delta_before")
    base = _array(base_candidate, np.float32, "base_candidate")
    comps = candidate_compartments
    _require(isinstance(comps, np.ndarray) and comps.ndim == 1 and comps.dtype.kind == "U",
             "candidate_compartments must be a Unicode vector")
    _require(len(delta) <= _INT32_MAX and len(sources) == len(delta) == len(base) == len(comps),
             "candidate columns have different lengths or exceed int32 capacity")
    _require(compartment in COMPARTMENTS and np.all(np.isin(comps, COMPARTMENTS)),
             "unrecognized exact frozen compartment")
    local = np.searchsorted(prepared.kc, sources)
    _require(np.all(local < len(prepared.kc)), "all candidate sources must be canonical KCs")
    _require(np.all(prepared.kc[local] == sources), "all candidate sources must be canonical KCs")
    d = _scalar(d, "D")
    _require(d in (0.0, 1.0), "D must be the fixed artificial scalar 0 or 1")
    _require(isinstance(frozen, (bool, np.bool_)), "frozen must be Boolean")
    eta = _scalar(eta_mv, "eta_mv")
    floor = _scalar(floor_fraction, "floor_fraction")
    decay_value = _scalar(decay, "decay")
    _require(0 < eta <= np.finfo(np.float32).max, "eta_mv must be positive finite float32 representable")
    _require(0 < floor <= 1, "floor_fraction must lie in (0,1]")
    _require(0 < decay_value <= 1, "decay must lie in (0,1]")
    _require(np.float32(eta) > 0 and np.float32(floor) > 0 and np.float32(decay_value) > 0,
             "rule scalar underflows float32")
    _require(np.all(np.isfinite(base)) and np.all(base > 0),
             "original base must be finite positive float32 mV")
    lower = np.float32(-floor) * base
    _require(np.all(np.isfinite(delta)) and np.all((delta >= lower) & (delta <= 0)),
             "delta violates original frozen floor")
    _require(not np.any((delta == 0) & np.signbit(delta)), "delta zero must be canonical +0")
    return dict(local=np.array(local, dtype=np.int32, order="C"),
                selected=np.asarray(comps == compartment, dtype=np.uint8),
                delta=np.array(delta, copy=True, order="C"),
                base=np.array(base, copy=True, order="C"), d=d, frozen=bool(frozen),
                eta=eta, floor=floor, decay=decay_value,
                negative_eta_times_d=np.float32(-eta * d), negative_floor=np.float32(-floor))


def _audit(prepared: PreparedEvents, work: dict, backend: str) -> dict[str, Any]:
    return dict(backend=backend, steps=prepared.steps, kc_count=len(prepared.kc),
                events=prepared.input_event_count, actual_kc_events=prepared.kc_event_count,
                command_count=len(prepared.command_steps), candidate_count=len(work["base"]),
                d_command=work["d"], frozen=work["frozen"], eta_mv=work["eta"],
                floor_fraction=work["floor"], decay_python_float=work["decay"],
                decay_float32_bits=int(np.asarray(np.float32(work["decay"])).view(np.uint32)),
                initial_trace="float32 canonical +0", sampling="after_current_tick_spike",
                update_sample="last command snapshot, or canonical +0 without commands",
                trace_tick_skipping=False, dense_tick_neuron_raster=False,
                original_source_executed=False, caller_must_verify_source_identity=True,
                source_identity_verified_by_this_module=False,
                complete_mapping_verified_by_this_module=False,
                d_semantics="artificial scheduled command; not measured DAN activity",
                actual_dan_activity_used_to_gate=False, brain_dynamics_executed=False,
                training_executed=False, weights_installed=False, learning_pass_claimed=False)


def cpu_reference(prepared: PreparedEvents, candidate_sources: np.ndarray,
                  candidate_compartments: np.ndarray, compartment: str,
                  delta_before: np.ndarray, base_candidate: np.ndarray, *, d: float,
                  frozen: bool, eta_mv: float, floor_fraction: float, decay: float) -> EligibilityResult:
    """Host arithmetic mirror for tests; frozen-source validation is external."""
    started = time.perf_counter()
    work = _validate_run(prepared, candidate_sources, candidate_compartments, compartment,
                         delta_before, base_candidate, d, frozen, eta_mv, floor_fraction, decay)
    nkc = len(prepared.kc)
    trace = np.zeros(nkc, np.float32)
    snapshots = np.empty((len(prepared.command_steps), nkc), np.float32)
    positions = prepared.indptr[:-1].copy()
    ends = prepared.indptr[1:]
    spike = np.zeros(nkc, np.float32)
    command = 0
    for tick in range(prepared.steps):
        spike.fill(0)
        active = np.flatnonzero(positions < ends)
        active = active[prepared.ticks[positions[active]] == tick]
        spike[active] = np.float32(1)
        positions[active] += 1
        np.multiply(trace, np.float32(work["decay"]), out=trace)
        np.add(trace, spike, out=trace)
        np.minimum(trace, np.float32(1), out=trace)
        if command < len(prepared.command_steps) and prepared.command_steps[command] == tick:
            snapshots[command] = trace
            command += 1
    sampled = snapshots[-1].copy() if len(snapshots) else np.zeros(nkc, np.float32)
    output = work["delta"].copy()
    if not work["frozen"]:
        update = np.zeros(len(output), np.float32)
        selected = work["selected"].astype(bool)
        if work["d"] != 0:
            update[selected] = work["negative_eta_times_d"] * sampled[work["local"][selected]]
        output = output + update
        output = np.maximum(output, work["negative_floor"] * work["base"]).astype(np.float32)
    return EligibilityResult(_owned(sampled), _owned(trace), _owned(snapshots), _owned(output),
                             dict(end_to_end_seconds=time.perf_counter() - started,
                                  preparation_seconds=prepared.preparation_seconds),
                             _audit(prepared, work, "numpy-host-math-mirror"))


class CudaEligibility:
    """CuPy RawKernel candidate with explicit round-to-nearest float32 math.

    Use a dedicated environment and set an absolute CUPY_CACHE_DIR outside any
    frozen tree before construction. The first run explicitly compiles/loads both
    kernels. cold_compile_seconds measures that first-use compile/cache-load,
    and does not prove an empty disk cache. Construction time is reported apart.
    run includes validation, transfers and output ownership in end-to-end wall.
    """

    def __init__(self, device_index: int = 0):
        started = time.perf_counter()
        self.device_index = _integer(device_index, "device_index", 0)
        cache_dir = os.environ.get("CUPY_CACHE_DIR", "")
        _require(bool(cache_dir) and Path(cache_dir).is_absolute(),
                 "set CUPY_CACHE_DIR to an absolute candidate cache directory outside frozen trees")
        try:
            import cupy as cp
        except ImportError as exc:
            raise RuntimeError("CuPy unavailable; use the dedicated CUDA candidate environment") from exc
        self.cp = cp
        self.device = cp.cuda.Device(self.device_index)
        with self.device:
            self.stream = cp.cuda.Stream(non_blocking=True)
            self.trace_kernel = cp.RawKernel(CUDA_SOURCE, "eligibility_trace", options=CUDA_OPTIONS)
            self.delta_kernel = cp.RawKernel(CUDA_SOURCE, "eligibility_delta", options=CUDA_OPTIONS)
            self.device.synchronize()
        self.initialization_seconds = time.perf_counter() - started
        self._compiled = False

    def run(self, prepared: PreparedEvents, candidate_sources: np.ndarray,
            candidate_compartments: np.ndarray, compartment: str,
            delta_before: np.ndarray, base_candidate: np.ndarray, *, d: float,
            frozen: bool, eta_mv: float, floor_fraction: float, decay: float,
            kernel_warmup: int = 0, kernel_repeats: int = 1) -> EligibilityResult:
        """Transfer one fixture and optionally benchmark resident kernel repeats.

        Each repeat starts from zero trace and the SAME delta_before; repeats
        never accumulate training updates. CUDA timing fields are medians across
        kernel_repeats. end_to_end_seconds includes all warmups and repeats, so
        use default counts for a single-call end-to-end benchmark.
        """
        started = time.perf_counter()
        work = _validate_run(prepared, candidate_sources, candidate_compartments, compartment,
                             delta_before, base_candidate, d, frozen, eta_mv, floor_fraction, decay)
        kernel_warmup = _integer(kernel_warmup, "kernel_warmup", 0)
        kernel_repeats = _integer(kernel_repeats, "kernel_repeats")
        validation_seconds = time.perf_counter() - started
        cp = self.cp
        first_run = not self._compiled
        compile_seconds = 0.0
        with self.device, self.stream:
            if first_run:
                compile_started = time.perf_counter()
                self.trace_kernel.compile()
                self.delta_kernel.compile()
                self.stream.synchronize()
                compile_seconds = time.perf_counter() - compile_started
                self._compiled = True
            transfer_started = time.perf_counter()
            ptr_gpu, ticks_gpu, commands_gpu, local_gpu, selected_gpu, delta_gpu, base_gpu = [
                cp.asarray(a) for a in (prepared.indptr, prepared.ticks, prepared.command_steps,
                                       work["local"], work["selected"], work["delta"], work["base"])]
            nkc, nedges = len(prepared.kc), len(work["delta"])
            snapshots_gpu = cp.empty((len(prepared.command_steps), nkc), cp.float32)
            final_gpu = cp.empty(nkc, cp.float32)
            sampled_gpu = cp.empty(nkc, cp.float32)
            output_gpu = cp.empty(nedges, cp.float32)
            self.stream.synchronize()
            upload_seconds = time.perf_counter() - transfer_started
            trace_args = (
                ptr_gpu, ticks_gpu, commands_gpu, np.int32(len(prepared.command_steps)),
                np.int32(nkc), np.int32(prepared.steps), np.float32(work["decay"]),
                snapshots_gpu, final_gpu, sampled_gpu)
            delta_args = (
                sampled_gpu, local_gpu, selected_gpu, delta_gpu, base_gpu,
                np.int32(nedges), work["negative_eta_times_d"], work["negative_floor"],
                np.int32(work["d"] != 0), np.int32(work["frozen"]), output_gpu)

            def launch_pair():
                self.trace_kernel(((nkc + 127) // 128,), (128,), trace_args)
                if nedges:
                    self.delta_kernel(((nedges + 255) // 256,), (256,), delta_args)

            for _ in range(kernel_warmup):
                launch_pair()
            self.stream.synchronize()
            resident_samples = []
            for _ in range(kernel_repeats):
                trace_start, trace_end, update_end = cp.cuda.Event(), cp.cuda.Event(), cp.cuda.Event()
                trace_start.record(self.stream)
                self.trace_kernel(((nkc + 127) // 128,), (128,), trace_args)
                trace_end.record(self.stream)
                if nedges:
                    self.delta_kernel(((nedges + 255) // 256,), (256,), delta_args)
                update_end.record(self.stream)
                update_end.synchronize()
                resident_samples.append(dict(
                    trace_kernel_seconds=cp.cuda.get_elapsed_time(trace_start, trace_end) / 1000.0,
                    update_kernel_seconds=cp.cuda.get_elapsed_time(trace_end, update_end) / 1000.0,
                    kernel_seconds=cp.cuda.get_elapsed_time(trace_start, update_end) / 1000.0))
            medians = {key: float(np.median([sample[key] for sample in resident_samples]))
                       for key in resident_samples[0]}
            download_started = time.perf_counter()
            sampled, final, snapshots, output = [
                _owned(cp.asnumpy(a, stream=self.stream, blocking=True))
                for a in (sampled_gpu, final_gpu, snapshots_gpu, output_gpu)]
            download_seconds = time.perf_counter() - download_started
        guard_started = time.perf_counter()
        for name, array in (("sampled_trace", sampled), ("final_trace", final), ("snapshots", snapshots)):
            if not (np.all(np.isfinite(array)) and np.all((array >= 0) & (array <= 1))):
                raise RuntimeError(f"CUDA output guard failed: invalid {name}")
        lower = work["negative_floor"] * work["base"]
        if not (np.all(np.isfinite(output)) and np.all((output >= lower) & (output <= 0))):
            raise RuntimeError("CUDA output guard failed: delta outside finite frozen floor")
        selected = work["selected"].astype(bool)
        if output[~selected].tobytes() != work["delta"][~selected].tobytes():
            raise RuntimeError("CUDA output guard failed: off-compartment delta changed")
        if (work["frozen"] or work["d"] == 0) and output.tobytes() != work["delta"].tobytes():
            raise RuntimeError("CUDA output guard failed: frozen/zero-D identity")
        guard_seconds = time.perf_counter() - guard_started
        timing = dict(first_run=first_run, initialization_seconds=self.initialization_seconds,
                      cold_compile_seconds=compile_seconds,
                      compile_timing_semantics="first-use compile/cache-load; cache coldness externally controlled",
                      preparation_seconds=prepared.preparation_seconds,
                      run_validation_seconds=validation_seconds,
                      host_to_device_and_allocation_seconds=upload_seconds,
                      **medians, device_to_host_seconds=download_seconds,
                      output_guard_seconds=guard_seconds,
                      resident_kernel_warmup=kernel_warmup,
                      resident_kernel_repeats=kernel_repeats,
                      resident_kernel_samples=resident_samples,
                      kernel_statistic="median of resident timed repeats",
                      end_to_end_semantics="all requested warmups/repeats and transfers; preparation excluded",
                      end_to_end_seconds=time.perf_counter() - started)
        audit = _audit(prepared, work, "cupy-rawkernel-independent-candidate")
        audit.update(numpy_version=str(np.__version__), cupy_version=str(cp.__version__),
                     device_index=self.device_index, cuda_options=list(CUDA_OPTIONS),
                     float32_separate_rounding_intrinsics=True, tf32_used=False,
                     numerical_equivalence_established_by_this_call=False)
        return EligibilityResult(sampled, final, snapshots, output, timing, audit)