// Per-millisecond LIF tick on CUDA, bit-exact with the CPU kernel `lif-1ms-f64-v2`.
//
// Mirrors steps 1-4 and 6 of `LifNetwork::step_one` in `src/lif.rs`: the noise kicks, the visual
// drive, the stimulation drive, the integrate-and-fire sweep, and spike propagation. Plasticity
// `observe`, the role tally and the rate EMAs stay on the host.
//
// Bit-exactness rules, all of which the host side re-checks against the CPU kernel:
//
//   * every intermediate is `double`; every store to a membrane rounds once with `(float)`, which
//     is IEEE round-to-nearest-even and therefore `Math.fround` / Rust's `as f32`;
//   * no FMA contraction (compiled with `--fmad=false`), no fast math, no reassociation;
//   * the only transcendental, `exp(-1/decayMs)`, is evaluated on the host at config time and
//     arrives as the already-frounded `decay` argument;
//   * propagation is a thread-per-target pull over a transposed CSR whose incoming edges are
//     ordered by base edge index, which is the source-major order the CPU push visits them in.
//     Skipping non-spiking sources leaves a subsequence in that same order, so each target sees
//     exactly the CPU's sequence of additions, with the f32 rounding and the per-edge floor clamp
//     at the same points.

#define LIF_BLOCK 256u

extern "C" {

// `Math.max`, which propagates NaN (unlike fmax) and returns +0 for max(-0, +0).
// Mirrors `jsmath::js_max`.
__device__ __forceinline__ double js_max(double a, double b) {
  if (!(a == a) || !(b == b)) {
    return __longlong_as_double(0x7ff8000000000000LL);
  }
  if (a > b) {
    return a;
  }
  if (b > a) {
    return b;
  }
  if (a == 0.0 && b == 0.0) {
    return (__double_as_longlong(a) >= 0LL) ? a : b;
  }
  return a;
}

// `lif::floor_clamp`: the common case first, `js_max` for the rest.
__device__ __forceinline__ double floor_clamp(double value, double floor) {
  return (value > floor) ? value : js_max(floor, value);
}

// Step 1. One thread per *distinct* neuron drawn this tick, applying that neuron's kicks in draw
// order. The host groups the 300 draws by neuron, so repeats (which the CPU applies in sequence
// through one f32 store each) stay sequential here too, and distinct neurons cannot interact.
__global__ void lif_noise(float *membrane, const unsigned int *neuron, const unsigned int *count,
                          unsigned int groups, double amount) {
  unsigned int group = blockIdx.x * blockDim.x + threadIdx.x;
  if (group >= groups) {
    return;
  }
  unsigned int index = neuron[group];
  unsigned int kicks = count[group];
  double value = (double)membrane[index];
  for (unsigned int kick = 0; kick < kicks; ++kick) {
    value = (double)(float)(value + amount);
  }
  membrane[index] = (float)value;
}

// Step 2. One thread per distinct visual target; `ptr`/`pos` list that neuron's columns in
// ascending column order, so a neuron that appears twice in `visual_indices` is still summed in
// the CPU's order.
__global__ void lif_visual(float *membrane, const unsigned int *neuron, const unsigned int *ptr,
                           const unsigned int *pos, const float *drive, unsigned int groups) {
  unsigned int group = blockIdx.x * blockDim.x + threadIdx.x;
  if (group >= groups) {
    return;
  }
  unsigned int index = neuron[group];
  double value = (double)membrane[index];
  for (unsigned int slot = ptr[group]; slot < ptr[group + 1]; ++slot) {
    value = (double)(float)(value + (double)drive[pos[slot]]);
  }
  membrane[index] = (float)value;
}

// Step 3. Same shape as the noise kicks: the drive is one constant, so a repeated target only
// needs its count.
__global__ void lif_stimulate(float *membrane, const unsigned int *neuron,
                              const unsigned int *count, unsigned int groups, double drive) {
  unsigned int group = blockIdx.x * blockDim.x + threadIdx.x;
  if (group >= groups) {
    return;
  }
  unsigned int index = neuron[group];
  unsigned int hits = count[group];
  double value = (double)membrane[index];
  for (unsigned int hit = 0; hit < hits; ++hit) {
    value = (double)(float)(value + drive);
  }
  membrane[index] = (float)value;
}

// Step 4. Elementwise integrate-and-fire, one thread per neuron, in the same `double` expression
// `lif::sweep_one` uses. Spikes leave a per-warp ballot word in `spiked_bits` (which propagation
// reads directly) plus a per-block count for the compaction scan; the ascending spike *list* is
// built by `lif_compact` from those bits.
__global__ void lif_sweep(float *membrane, unsigned char *refractory, const float *baseline,
                          unsigned int neurons, double decay, double threshold,
                          unsigned int refractory_reset, unsigned int *spiked_bits,
                          unsigned int *block_counts) {
  unsigned int index = blockIdx.x * blockDim.x + threadIdx.x;
  bool spiked = false;
  if (index < neurons) {
    unsigned char left = refractory[index];
    if (left > 0) {
      // A refractory neuron still leaks but receives no baseline drive.
      refractory[index] = (unsigned char)(left - 1);
      membrane[index] = (float)((double)membrane[index] * decay);
    } else {
      double voltage = (double)membrane[index] * decay + (double)baseline[index];
      if (voltage >= threshold) {
        // A spiking neuron resets to exactly 0 rather than subtracting the threshold.
        membrane[index] = 0.0f;
        refractory[index] = (unsigned char)refractory_reset;
        spiked = true;
      } else {
        membrane[index] = (float)voltage;
      }
    }
  }

  unsigned int mask = __ballot_sync(0xffffffffu, spiked);
  unsigned int lane = threadIdx.x & 31u;
  unsigned int warp = threadIdx.x >> 5;
  __shared__ unsigned int per_warp[LIF_BLOCK / 32u];
  if (lane == 0u) {
    unsigned int word = (blockIdx.x * blockDim.x + warp * 32u) >> 5;
    if (word < ((neurons + 31u) >> 5)) {
      spiked_bits[word] = mask;
    }
    per_warp[warp] = (unsigned int)__popc(mask);
  }
  __syncthreads();
  if (threadIdx.x == 0u) {
    unsigned int total = 0u;
    for (unsigned int slot = 0u; slot < (blockDim.x >> 5); ++slot) {
      total += per_warp[slot];
    }
    block_counts[blockIdx.x] = total;
  }
}

// Exclusive scan of the per-block spike counts, in one block, so `lif_compact` can write straight
// into the batch's flat spike buffer. Also records this tick's spike count and advances the shared
// write cursor, which keeps the whole batch's spike lists packed for a single download.
__global__ void lif_scan(const unsigned int *block_counts, unsigned int *block_offsets,
                         unsigned int blocks, unsigned int *tick_counts, unsigned int *tick_base,
                         unsigned int tick, unsigned int *cursor) {
  extern __shared__ unsigned int partial[];
  unsigned int thread = threadIdx.x;
  unsigned int threads = blockDim.x;
  unsigned int chunk = (blocks + threads - 1u) / threads;
  unsigned int from = thread * chunk;
  unsigned int to = from + chunk;
  if (from > blocks) {
    from = blocks;
  }
  if (to > blocks) {
    to = blocks;
  }

  unsigned int sum = 0u;
  for (unsigned int block = from; block < to; ++block) {
    sum += block_counts[block];
  }
  partial[thread] = sum;
  __syncthreads();
  for (unsigned int offset = 1u; offset < threads; offset <<= 1) {
    unsigned int carry = (thread >= offset) ? partial[thread - offset] : 0u;
    __syncthreads();
    partial[thread] += carry;
    __syncthreads();
  }

  unsigned int base = *cursor;
  unsigned int running = base + (partial[thread] - sum);
  for (unsigned int block = from; block < to; ++block) {
    block_offsets[block] = running;
    running += block_counts[block];
  }
  __syncthreads();
  if (thread == threads - 1u) {
    tick_counts[tick] = partial[threads - 1u];
    tick_base[tick] = base;
    *cursor = base + partial[threads - 1u];
  }
}

// The ascending spike list. Block offsets ascend with the neuron index and the rank inside a block
// is a ballot prefix, so the list is exactly the sequential kernel's: `observe` and the role tally
// consume it in this order.
__global__ void lif_compact(const unsigned int *spiked_bits, const unsigned int *block_offsets,
                            unsigned int neurons, unsigned int *spikes) {
  unsigned int index = blockIdx.x * blockDim.x + threadIdx.x;
  unsigned int lane = threadIdx.x & 31u;
  unsigned int warp = threadIdx.x >> 5;
  unsigned int word = (blockIdx.x * blockDim.x + warp * 32u) >> 5;
  unsigned int mask = (word < ((neurons + 31u) >> 5)) ? spiked_bits[word] : 0u;

  __shared__ unsigned int per_warp[LIF_BLOCK / 32u];
  if (lane == 0u) {
    per_warp[warp] = (unsigned int)__popc(mask);
  }
  __syncthreads();

  if (index < neurons && ((mask >> lane) & 1u) != 0u) {
    unsigned int rank = (unsigned int)__popc(mask & ((1u << lane) - 1u));
    for (unsigned int slot = 0u; slot < warp; ++slot) {
      rank += per_warp[slot];
    }
    spikes[block_offsets[blockIdx.x] + rank] = index;
  }
}

// Step 6, part 1 of 4: how many of this tick's edges land on each target.
//
// The pull formulation this replaced — a warp per target over a transposed CSR, skipping sources
// that did not spike — is bit-exact and was measured at 320 us per tick, because it has to read
// all 2,700,513 incoming edges every tick while the CPU push reads only the ~1.2 % that leave a
// spiking neuron (about 33,000 edges). That asymmetry, not the arithmetic, is what made the GPU
// slower than four CPU cores. So: push the spiking rows into per-target buckets, then let one
// thread per target sum its own bucket in the CPU's order.
//
// Bucket *placement* uses atomics and is therefore in an arbitrary order; `lif_bucket_apply`
// restores the order by sorting each bucket on the base edge index, which is exactly the CPU's
// order (the sequential kernel walks spiking sources ascending and each row ascending, and the
// base edge index is source-major then slot). Buckets hold a handful of entries, so the sort is
// free and determinism costs nothing.
//
// A grid-stride loop over the spiking sources, one warp each, so the row reads are coalesced and
// the grid does not have to be sized for the worst case.
__global__ void lif_bucket_count(const unsigned int *__restrict__ spikes,
                                 const unsigned int *__restrict__ tick_counts,
                                 const unsigned int *__restrict__ tick_base, unsigned int tick,
                                 const unsigned int *__restrict__ indptr,
                                 const unsigned int *__restrict__ targets, unsigned int *count) {
  unsigned int spiking = tick_counts[tick];
  unsigned int base = tick_base[tick];
  unsigned int warps = (gridDim.x * blockDim.x) >> 5;
  unsigned int warp = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
  unsigned int lane = threadIdx.x & 31u;
  for (unsigned int index = warp; index < spiking; index += warps) {
    unsigned int source = spikes[base + index];
    unsigned int from = indptr[source];
    unsigned int to = indptr[source + 1u];
    for (unsigned int edge = from + lane; edge < to; edge += 32u) {
      atomicAdd(&count[targets[edge]], 1u);
    }
  }
}

// Step 6, part 2 of 4: the same walk again, writing each edge into its target's bucket.
//
// A bucket entry is one `unsigned long long`: the base edge index in the top 32 bits, the edge's
// `short` weight in the bottom 16, and the plastic gain slot plus one (0 for the immutable
// majority) in bits 16-30. The edge index being the most significant field means sorting the
// entries as integers sorts them by edge index, which is the CPU's order.
//
// Resolving the weight and the gain slot *here* rather than in `lif_bucket_apply` is worth 80 us
// per tick: here the reads follow a CSR row, so a warp's 32 lanes read consecutive weights and
// share a plastic-bitset word; there they would be 97,000 scattered reads by whichever thread owns
// the target. The arithmetic is unaffected — the same `short` and the same `f32` gain reach the
// same f64 expression.
__global__ void lif_bucket_place(const unsigned int *__restrict__ spikes,
                                 const unsigned int *__restrict__ tick_counts,
                                 const unsigned int *__restrict__ tick_base, unsigned int tick,
                                 const unsigned int *__restrict__ indptr,
                                 const unsigned int *__restrict__ targets,
                                 const short *__restrict__ weights,
                                 const unsigned long long *__restrict__ plastic_words,
                                 const unsigned int *__restrict__ plastic_prefix,
                                 unsigned int *cursor, unsigned long long *bucket) {
  unsigned int spiking = tick_counts[tick];
  unsigned int base = tick_base[tick];
  unsigned int warps = (gridDim.x * blockDim.x) >> 5;
  unsigned int warp = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
  unsigned int lane = threadIdx.x & 31u;
  for (unsigned int index = warp; index < spiking; index += warps) {
    unsigned int source = spikes[base + index];
    unsigned int from = indptr[source];
    unsigned int to = indptr[source + 1u];
    for (unsigned int edge = from + lane; edge < to; edge += 32u) {
      unsigned long long word = plastic_words[edge >> 6];
      unsigned int slot = 0u;
      if (((word >> (edge & 63u)) & 1ULL) != 0ULL) {
        unsigned long long below = word & ((1ULL << (edge & 63u)) - 1ULL);
        slot = plastic_prefix[edge >> 6] + (unsigned int)__popcll(below) + 1u;
      }
      unsigned long long entry = ((unsigned long long)edge << 32) |
                                 (unsigned long long)((unsigned int)(unsigned short)weights[edge]) |
                                 ((unsigned long long)slot << 16);
      bucket[atomicAdd(&cursor[targets[edge]], 1u)] = entry;
    }
  }
}

// Step 6, part 3 of 4: an exclusive scan of the bucket sizes, in two kernels plus the small
// single-block scan below. `offset` holds the block-local exclusive scan on the way in and the
// global one on the way out.
__global__ void lif_bucket_scan(const unsigned int *__restrict__ count, unsigned int *offset,
                                unsigned int *block_sums, unsigned int neurons) {
  __shared__ unsigned int partial[LIF_BLOCK];
  unsigned int index = blockIdx.x * blockDim.x + threadIdx.x;
  unsigned int thread = threadIdx.x;
  unsigned int mine = (index < neurons) ? count[index] : 0u;
  partial[thread] = mine;
  __syncthreads();
  for (unsigned int step = 1u; step < blockDim.x; step <<= 1) {
    unsigned int carry = (thread >= step) ? partial[thread - step] : 0u;
    __syncthreads();
    partial[thread] += carry;
    __syncthreads();
  }
  if (index < neurons) {
    offset[index] = partial[thread] - mine;
  }
  if (thread == blockDim.x - 1u) {
    block_sums[blockIdx.x] = partial[thread];
  }
}

// Exclusive scan of up to `blockDim.x` * chunk values, in one block. Used for the bucket scan's
// block sums.
__global__ void lif_exscan(const unsigned int *__restrict__ input, unsigned int *output,
                           unsigned int count) {
  extern __shared__ unsigned int partial[];
  unsigned int thread = threadIdx.x;
  unsigned int threads = blockDim.x;
  unsigned int chunk = (count + threads - 1u) / threads;
  unsigned int from = thread * chunk;
  unsigned int to = from + chunk;
  if (from > count) {
    from = count;
  }
  if (to > count) {
    to = count;
  }
  unsigned int sum = 0u;
  for (unsigned int index = from; index < to; ++index) {
    sum += input[index];
  }
  partial[thread] = sum;
  __syncthreads();
  for (unsigned int step = 1u; step < threads; step <<= 1) {
    unsigned int carry = (thread >= step) ? partial[thread - step] : 0u;
    __syncthreads();
    partial[thread] += carry;
    __syncthreads();
  }
  unsigned int running = partial[thread] - sum;
  for (unsigned int index = from; index < to; ++index) {
    output[index] = running;
    running += input[index];
  }
}

// Add each block's base to its block-local scan, and seed the placement cursor.
__global__ void lif_bucket_rebase(unsigned int *offset, unsigned int *cursor,
                                  const unsigned int *__restrict__ block_bases,
                                  unsigned int neurons) {
  unsigned int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index >= neurons) {
    return;
  }
  unsigned int value = offset[index] + block_bases[blockIdx.x];
  offset[index] = value;
  cursor[index] = value;
}

// Step 6, part 4 of 4: one thread per target sums its own bucket, in the CPU's order.
//
// Every target's membrane depends on nothing but its own previous value and its own sequence of
// addends, so this is where bit-exactness is won or lost: the bucket is sorted on the base edge
// index, each addend is the same `weight * gain * synapse_scale` product evaluated in the same
// order, and the f32 rounding and the per-edge floor clamp happen at the same points.
__global__ void lif_bucket_apply(float *membrane, const unsigned int *__restrict__ count,
                                 const unsigned int *__restrict__ offset,
                                 unsigned long long *bucket, const float *__restrict__ gains,
                                 unsigned int neurons, double synapse_scale,
                                 double membrane_floor) {
  unsigned int target = blockIdx.x * blockDim.x + threadIdx.x;
  if (target >= neurons) {
    return;
  }
  unsigned int entries = count[target];
  if (entries == 0u) {
    return;
  }
  unsigned int from = offset[target];

  // Sort the bucket, whose placement order is whatever the atomics produced.
  //
  // Shell sort with Ciura's gaps, not plain insertion sort. Buckets average well under one entry,
  // so the gap loop degenerates to insertion sort for almost every target — but the neuron with
  // the maximum in-degree (5,080 on fafb-v783) collects a couple of hundred entries in a tick, one
  // thread owns all of them, and that thread's k^2 was measured as 100 us of the tick. Any correct
  // sort preserves bit-exactness, because the keys are distinct edge indices.
  const unsigned int gaps[8] = {701u, 301u, 132u, 57u, 23u, 10u, 4u, 1u};
  for (unsigned int step = 0u; step < 8u; ++step) {
    unsigned int gap = gaps[step];
    if (gap >= entries) {
      continue;
    }
    for (unsigned int index = gap; index < entries; ++index) {
      unsigned long long key = bucket[from + index];
      unsigned int at = index;
      while (at >= gap && bucket[from + at - gap] > key) {
        bucket[from + at] = bucket[from + at - gap];
        at -= gap;
      }
      bucket[from + at] = key;
    }
  }

  double value = (double)membrane[target];
  for (unsigned int index = 0u; index < entries; ++index) {
    unsigned long long entry = bucket[from + index];
    short weight = (short)(unsigned short)(entry & 0xffffULL);
    unsigned int slot = (unsigned int)((entry >> 16) & 0x7fffULL);
    double gain = (slot == 0u) ? 1.0 : (double)gains[slot - 1u];
    // The floor is applied per edge, not once per tick: the clamp order is part of the numerics.
    double next = value + (double)weight * gain * synapse_scale;
    value = (double)(float)floor_clamp(next, membrane_floor);
  }
  membrane[target] = (float)value;
}

}  // extern "C"