#include <torch/extension.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cstdint>

namespace {

__global__ void state_kernel(float* v, float* g, int* refractory,
                             float* arrival, unsigned char* spikes,
                             int* spike_counts, int n, float g_decay,
                             float alpha, float rest, float reset,
                             float threshold, int refractory_steps) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= n) return;
  // Phase 4B executes mul_ and add_ as separate kernels: prevent contraction.
  float gi = __fadd_rn(__fmul_rn(g[i], g_decay), arrival[i]);
  g[i] = gi;
  arrival[i] = 0.0f;
  int r = refractory[i];
  r = r > 0 ? r - 1 : 0;
  refractory[i] = r;
  float vi = v[i];
  bool available = r == 0;
  if (available) {
    // Match the separate float32 operations used by the eager reference.
    float leak = __fsub_rn(vi, rest);
    float drive = __fsub_rn(gi, leak);
    vi = __fadd_rn(vi, __fmul_rn(alpha, drive));
  } else {
    vi = reset;
  }
  bool fired = available && vi >= threshold;
  if (fired) {
    vi = reset;
    refractory[i] = refractory_steps;
    spike_counts[i] += 1;
  }
  v[i] = vi;
  spikes[i] = fired ? 1 : 0;
}

__global__ void force_kernel(float* v, int* refractory,
                             unsigned char* spikes, int* spike_counts,
                             const int64_t* targets,
                             const unsigned char* forced, int count,
                             float reset, int refractory_steps) {
  int j = blockIdx.x * blockDim.x + threadIdx.x;
  if (j >= count || !forced[j]) return;
  int64_t i = targets[j];
  // A natural spike already recorded at this target represents the same event.
  if (spikes[i]) return;
  if (refractory[i] == 0) {
    spikes[i] = 1;
    v[i] = reset;
    refractory[i] = refractory_steps;
    spike_counts[i] += 1;
  }
}

__global__ void fanout_kernel(const unsigned char* spikes,
                              const int64_t* indptr, const int* indices,
                              const float* weights, float* destination,
                              int n) {
  int src = blockIdx.x * blockDim.x + threadIdx.x;
  if (src >= n || !spikes[src]) return;
  int64_t start = indptr[src];
  int64_t stop = indptr[src + 1];
  for (int64_t e = start; e < stop; ++e) {
    atomicAdd(destination + indices[e], weights[e]);
  }
}

__global__ void fill_edges_kernel(const int64_t* indptr,
                                  const int64_t* sources,
                                  const int64_t* output_offsets,
                                  int64_t* output, int source_count) {
  int j = blockIdx.x * blockDim.x + threadIdx.x;
  if (j >= source_count) return;
  int64_t src = sources[j];
  int64_t begin = indptr[src];
  int64_t degree = indptr[src + 1] - begin;
  int64_t out = output_offsets[j];
  for (int64_t k = 0; k < degree; ++k) output[out + k] = begin + k;
}

void check_cuda(const torch::Tensor& x) {
  TORCH_CHECK(x.is_cuda(), "all tensors must be CUDA tensors");
  TORCH_CHECK(x.is_contiguous(), "all tensors must be contiguous");
}
}  // namespace

void state_step_cuda(torch::Tensor v, torch::Tensor g, torch::Tensor refractory,
                     torch::Tensor arrival, torch::Tensor spikes,
                     torch::Tensor spike_counts, double g_decay, double alpha,
                     double rest, double reset, double threshold,
                     int64_t refractory_steps) {
  check_cuda(v); check_cuda(g); check_cuda(refractory); check_cuda(arrival);
  check_cuda(spikes); check_cuda(spike_counts);
  TORCH_CHECK(v.scalar_type() == torch::kFloat32 && g.scalar_type() == torch::kFloat32,
              "v/g must be float32");
  TORCH_CHECK(refractory.scalar_type() == torch::kInt32 && spike_counts.scalar_type() == torch::kInt32,
              "refractory/counts must be int32");
  TORCH_CHECK(spikes.scalar_type() == torch::kUInt8, "spikes must be uint8");
  int n = v.numel();
  state_kernel<<<(n + 255) / 256, 256>>>(
      v.data_ptr<float>(), g.data_ptr<float>(), refractory.data_ptr<int>(),
      arrival.data_ptr<float>(), spikes.data_ptr<unsigned char>(),
      spike_counts.data_ptr<int>(), n, (float)g_decay, (float)alpha,
      (float)rest, (float)reset, (float)threshold, (int)refractory_steps);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void force_step_cuda(torch::Tensor v, torch::Tensor refractory,
                     torch::Tensor spikes, torch::Tensor spike_counts,
                     torch::Tensor targets, torch::Tensor forced,
                     double reset, int64_t refractory_steps) {
  check_cuda(targets); check_cuda(forced);
  int count = targets.numel();
  force_kernel<<<(count + 255) / 256, 256>>>(
      v.data_ptr<float>(), refractory.data_ptr<int>(),
      spikes.data_ptr<unsigned char>(), spike_counts.data_ptr<int>(),
      targets.data_ptr<int64_t>(), forced.data_ptr<unsigned char>(), count,
      (float)reset, (int)refractory_steps);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void fanout_cuda(torch::Tensor spikes, torch::Tensor indptr,
                 torch::Tensor indices, torch::Tensor weights,
                 torch::Tensor destination) {
  check_cuda(spikes); check_cuda(indptr); check_cuda(indices);
  check_cuda(weights); check_cuda(destination);
  int n = spikes.numel();
  fanout_kernel<<<(n + 255) / 256, 256>>>(
      spikes.data_ptr<unsigned char>(), indptr.data_ptr<int64_t>(),
      indices.data_ptr<int>(), weights.data_ptr<float>(),
      destination.data_ptr<float>(), n);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void fill_edges_cuda(torch::Tensor indptr, torch::Tensor sources,
                     torch::Tensor output_offsets, torch::Tensor output) {
  check_cuda(indptr); check_cuda(sources); check_cuda(output_offsets); check_cuda(output);
  TORCH_CHECK(indptr.scalar_type() == torch::kInt64 && sources.scalar_type() == torch::kInt64 &&
              output_offsets.scalar_type() == torch::kInt64 && output.scalar_type() == torch::kInt64,
              "fill_edges tensors must be int64");
  int count = sources.numel();
  fill_edges_kernel<<<(count + 255) / 256, 256>>>(
      indptr.data_ptr<int64_t>(), sources.data_ptr<int64_t>(),
      output_offsets.data_ptr<int64_t>(), output.data_ptr<int64_t>(), count);
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}