"""CPU-only, target-local shadow of the frozen author LIF + legacy APL schedule.

Targets generate their own threshold events.  Recorded whole-graph events are
used ONLY as upstream, fixed-delay inputs; recorded target spikes are never
forced into the target state.  APL h is an observed external component trace.
This is a conditional shadow reconstruction, not independent whole-brain or
independent APL validation.  Nothing in this module imports Brian or Torch,
runs a scientific trial, reads a file, or writes a file.
"""
from __future__ import annotations

from collections.abc import Mapping

import numpy as np


N_GLOBAL = 166700
DELAY_STEPS = 18
REF_STEPS = 22
COEFFICIENT_KEYS = ("gd", "c0", "cg", "cv", "threshold", "reset", "seed_volt")


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


def _vector(value, dtype, name, size=None, finite=False):
    value = np.asarray(value)
    _require(value.dtype == np.dtype(dtype) and value.ndim == 1, f"{name} must be a one-dimensional {dtype} array")
    _require(size is None or len(value) == size, f"{name} has the wrong length")
    if finite:
        _require(np.isfinite(value).all(), f"{name} contains a nonfinite value")
    return value


def _coefficients(coeff):
    _require(isinstance(coeff, Mapping), "coeff must be a mapping")
    result = {}
    for name in COEFFICIENT_KEYS:
        _require(name in coeff, f"Missing coefficient: {name}")
        value = coeff[name]
        _require(not isinstance(value, (bool, np.bool_)) and isinstance(value, (int, float, np.integer, np.float64)),
                 f"{name} must be a double-precision scalar, not a float32 or nonnumeric value")
        result[name] = np.float64(value)
        _require(np.isfinite(result[name]), f"Nonfinite coefficient: {name}")
    _require(result["seed_volt"] > result["threshold"], "Common-request seed voltage must exceed the strict threshold")
    return result


class IncomingIndex:
    """Reusable immutable target-filtered anatomical edges in original order.

    src_global:int32, dst_local:int32, edge_id:int64, weights:float64 (SI volts).
    All vectors have the same length.  Sources and original global edge IDs
    must retain canonical increasing source/edge order.  Edges are NOT removed
    because their weight is zero, nor reordered by target.  The caller must
    supply the already-audited effective weights (all fast APL source edges
    zeroed in a separate copy); this object never edits its input arrays.
    """

    def __init__(self, src_global, dst_local, edge_id, weights, target_count):
        _require(isinstance(target_count, (int, np.integer)) and not isinstance(target_count, (bool, np.bool_)) and
                 0 < target_count <= N_GLOBAL, "target_count must be a positive in-range integer")
        self.target_count = int(target_count)
        source = _vector(src_global, "int32", "src_global")
        count = len(source)
        destination = _vector(dst_local, "int32", "dst_local", count)
        ids = _vector(edge_id, "int64", "edge_id", count)
        weight = _vector(weights, "float64", "weights", count, finite=True)
        _require(np.all((source >= 0) & (source < N_GLOBAL)), "An incoming source is outside the global neuron range")
        _require(np.all((destination >= 0) & (destination < target_count)), "Incoming destination is outside the target-local range")
        _require(np.all(ids >= 0) and np.all(ids[1:] > ids[:-1]), "Original edge IDs must be unique and canonically increasing")
        _require(np.all(source[1:] >= source[:-1]), "Incoming sources must remain in original canonical order")
        self.src_global = np.array(source, dtype=np.int32, copy=True, order="C")
        self.dst_local = np.array(destination, dtype=np.int32, copy=True, order="C")
        self.edge_id = np.array(ids, dtype=np.int64, copy=True, order="C")
        self.weights = np.array(weight, dtype=np.float64, copy=True, order="C")
        self.sources, starts = np.unique(self.src_global, return_index=True)
        self.starts = np.asarray(starts, dtype=np.int64)
        self.stops = np.concatenate((self.starts[1:], np.array([count], np.int64))) if count else np.empty(0, np.int64)
        for value in (self.src_global, self.dst_local, self.edge_id, self.weights, self.sources, self.starts, self.stops):
            value.setflags(write=False)

    def due(self, global_spikes):
        """Return local edge rows in original queue order; no target sort/sum."""
        fired = _vector(global_spikes, "int32", "global_spikes")
        _require(np.all((fired >= 0) & (fired < N_GLOBAL)) and np.all(fired[1:] > fired[:-1]),
                 "Presynaptic spikes must be sorted, unique, in-range globals")
        if not len(fired) or not len(self.sources):
            return np.empty(0, np.int64)
        positions = np.searchsorted(self.sources, fired)
        present = positions < len(self.sources)
        present[present] &= self.sources[positions[present]] == fired[present]
        positions = positions[present]
        if not len(positions):
            return np.empty(0, np.int64)
        # Each presynaptic group occupies a contiguous source-ordered block in
        # this filtered graph.  Concatenation preserves every target's original
        # per-edge accumulation order, including zero-weight events.
        return np.concatenate([np.arange(self.starts[pos], self.stops[pos], dtype=np.int64) for pos in positions])


class _EventIndex:
    def __init__(self, record, steps, name):
        _require(isinstance(record, Mapping), f"{name} must be a mapping")
        time_keys = [key for key in ("tick", "step") if key in record]
        _require(len(time_keys) == 1 and "neuron_idx" in record, f"{name} needs exactly one tick/step vector and neuron_idx")
        ticks = _vector(record[time_keys[0]], "int32", name + ".tick")
        nodes = _vector(record["neuron_idx"], "int32", name + ".neuron_idx", len(ticks))
        _require(np.all((ticks >= 0) & (ticks < steps)), f"{name} has an event outside the trial")
        _require(np.all((nodes >= 0) & (nodes < N_GLOBAL)), f"{name} contains an out-of-range global neuron")
        _require(np.all((ticks[1:] > ticks[:-1]) | ((ticks[1:] == ticks[:-1]) & (nodes[1:] > nodes[:-1]))),
                 f"{name} must be sorted and unique by tick/global neuron")
        self.nodes = np.array(nodes, copy=True, order="C")
        counts = np.bincount(ticks, minlength=steps)
        self.pointer = np.concatenate((np.zeros(1, np.int64), np.cumsum(counts, dtype=np.int64)))
        self.nodes.setflags(write=False)
        self.pointer.setflags(write=False)

    def at(self, tick):
        if tick < 0:
            return np.empty(0, np.int32)
        return self.nodes[self.pointer[tick]:self.pointer[tick + 1]]


def _state_step(v, g, ref, request, coeff, capture):
    """Separate NumPy array operations correspond to explicit CUDA RN ops."""
    next_ref = np.maximum(ref - np.int32(1), np.int32(0))
    available = next_ref == 0
    next_v, next_g = v.copy(), g.copy()
    oldg, oldv = g[available], v[available]
    next_g[available] = np.multiply(coeff["gd"], oldg)
    driven = np.multiply(coeff["cg"], oldg)
    leaked = np.multiply(coeff["cv"], oldv)
    combined = np.add(driven, leaked)
    next_v[available] = np.add(coeff["c0"], combined)
    result = {}
    if capture:
        result.update(old_v=v.copy(), old_g=g.copy(), ref_before=ref.copy(),
                      available=available.copy(), request=request.copy(),
                      request_accepted=(request & available), request_rejected=(request & ~available),
                      v_integrated=next_v.copy(), g_decayed=next_g.copy())
    next_v[request & available] = coeff["seed_volt"]
    if capture:
        result["v_threshold"] = next_v.copy()
    fired = available & (next_v > coeff["threshold"])
    next_v[fired] = coeff["reset"]
    next_g[fired] = np.float64(0.0)
    next_ref[fired] = np.int32(REF_STEPS)
    _require(np.isfinite(next_v).all() and np.isfinite(next_g).all(), "State step produced a nonfinite value")
    result.update(v=next_v, g=next_g, ref=next_ref, fired=fired)
    if capture:
        result.update(v_after_reset=next_v.copy(), g_after_reset=next_g.copy(), ref_after_reset=next_ref.copy())
    return result


def state_step(v, g, ref, request, coeff, *, capture=True):
    """One group/request/threshold/reset step; no arrivals or APL yet.

    Inputs are never modified.  v/g are float64 SI, ref is int32 in [0,22],
    request is a bool vector.  With capture=True this returns the corresponding
    pre-threshold/intermediate callback fields as well as v/g/ref/fired.
    """
    voltage = _vector(v, "float64", "v", finite=True)
    n = len(voltage)
    conductance = _vector(g, "float64", "g", n, finite=True)
    refractory = _vector(ref, "int32", "ref", n)
    requested = _vector(request, "bool", "request", n)
    _require(n > 0 and np.all((refractory >= 0) & (refractory <= REF_STEPS)), "Invalid target count/refractory state")
    _require(isinstance(capture, bool), "capture must be Boolean")
    return _state_step(voltage, conductance, refractory, requested, _coefficients(coeff), capture)


def replay(initial_v, initial_g, coeff, steps, target_globals, requests, events,
           incoming: IncomingIndex, h32, callback=None):
    """Reconstruct targets with bounded working state and sparse output events.

    initial_v/initial_g: float64 SI vectors, initially nonrefractory.
    target_globals: sorted unique int32 target IDs.
    requests/events: mappings with tick (or step) and neuron_idx int32 vectors;
      requests outside the target subset are ignored, not reassigned.
    h32: float32 (steps, target_count), ALREADY mapped to target_globals, in mV.
      Full 4064-KC input is valid when the target list is those same 4064 KCs.
    callback(tick, stage): optional copies of old states, available/request
      masks, integrated/threshold voltages, reset state, canonical due edges,
      accepted arrivals, pre-APL states, h32/delta_si, and final v/g/ref.  It may
      raise to stop at the first caller-detected mismatch.  Modifying callback
      arrays cannot affect this replay or its input arrays.

    No dense history is stored.  Memory is O(incoming edges + logged events +
    steps + targets + current due edges + predicted sparse events), excluding
    caller-owned h32 and anything the callback elects to retain.
    """
    _require(isinstance(steps, (int, np.integer)) and not isinstance(steps, (bool, np.bool_)) and
             0 < steps <= np.iinfo(np.int32).max, "steps must be a positive int32-range integer")
    steps = int(steps)
    targets = _vector(target_globals, "int32", "target_globals")
    n = len(targets)
    _require(0 < n <= N_GLOBAL and np.all((targets >= 0) & (targets < N_GLOBAL)) and np.all(targets[1:] > targets[:-1]),
             "Targets must be sorted unique in-range globals")
    targets = targets.copy()
    v = _vector(initial_v, "float64", "initial_v", n, finite=True).copy()
    g = _vector(initial_g, "float64", "initial_g", n, finite=True).copy()
    constants = _coefficients(coeff)
    _require(isinstance(incoming, IncomingIndex) and incoming.target_count == n, "Incoming index target count mismatch")
    _require(callback is None or callable(callback), "callback must be callable or None")
    inhibition = np.asarray(h32)
    _require(inhibition.dtype == np.dtype("float32") and inhibition.shape == (steps, n), "h32 must be float32 (steps, targets), already target-mapped")
    for tick in range(steps):
        _require(np.isfinite(inhibition[tick]).all() and np.all(inhibition[tick] >= 0), "APL inhibition must be finite and nonnegative")
    request_index = _EventIndex(requests, steps, "requests")
    event_index = _EventIndex(events, steps, "events")
    ref = np.zeros(n, np.int32)
    counts = np.zeros(n, np.int32)
    first_spike = np.full(n, -1, np.int32)
    predicted_steps, predicted_nodes = [], []
    capture = callback is not None
    for tick in range(steps):
        requested_globals = request_index.at(tick)
        local = np.searchsorted(targets, requested_globals)
        present = local < n
        present[present] &= targets[local[present]] == requested_globals[present]
        request = np.zeros(n, dtype=bool)
        request[local[present]] = True
        stage = _state_step(v, g, ref, request, constants, capture)
        v, g, ref, fired = stage["v"], stage["g"], stage["ref"], stage["fired"]
        counts += fired.astype(np.int32)
        first_spike[fired & (first_spike < 0)] = np.int32(tick)
        fired_local = np.flatnonzero(fired)
        if len(fired_local):
            predicted_steps.append(np.full(len(fired_local), tick, np.int32))
            predicted_nodes.append(targets[fired_local].copy())
        due = incoming.due(event_index.at(tick - DELAY_STEPS))
        accepted = ref[incoming.dst_local[due]] == 0
        if len(due):
            accepted_edges = due[accepted]
            # add.at is deliberately unbuffered: start at existing g and add
            # each accepted edge in original order, NEVER sum weights first.
            np.add.at(g, incoming.dst_local[accepted_edges], incoming.weights[accepted_edges])
        if capture:
            stage.update(due_edge_indices=due.copy(), due_edge_id=incoming.edge_id[due].copy(),
                         arrival_accepted=accepted.copy(), v_pre_apl=v.copy(), g_pre_apl=g.copy(), ref_pre_apl=ref.copy())
        delta_si = inhibition[tick].astype(np.float64) * np.float64(.001)
        g = np.subtract(g, delta_si)
        _require(np.isfinite(v).all() and np.isfinite(g).all(), "Arrival/APL stage produced a nonfinite value")
        if capture:
            stage.update(h32=inhibition[tick].copy(), delta_si=delta_si.copy(), v=v.copy(), g=g.copy(), ref=ref.copy(), fired=fired.copy())
            callback(tick, stage)
    return {"v": v.copy(), "g": g.copy(), "ref": ref.copy(), "counts": counts,
            "first_spike": first_spike, "target_globals": targets,
            "step": np.concatenate(predicted_steps) if predicted_steps else np.empty(0, np.int32),
            "neuron_idx": np.concatenate(predicted_nodes) if predicted_nodes else np.empty(0, np.int32)}