#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
#include <cstdint>
#include <limits>

namespace {
constexpr int BLOCK = 128;

void check_tensor(const torch::Tensor& x, const torch::Tensor& like,
                  at::ScalarType dtype, const char* name) {
  TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.dim() == 1,
              name, " must be a contiguous CUDA vector");
  TORCH_CHECK(x.device() == like.device(), name, " is on a different CUDA device");
  TORCH_CHECK(x.scalar_type() == dtype, name, " has the wrong dtype");
}
void check_state(const torch::Tensor& g, const torch::Tensor& ref) {
  TORCH_CHECK(g.scalar_type() == torch::kFloat32 || g.scalar_type() == torch::kFloat64,
              "state must be float32 or float64");
  check_tensor(g, g, g.scalar_type(), "g");
  check_tensor(ref, g, torch::kInt32, "ref");
  TORCH_CHECK(g.numel() == ref.numel(), "g/ref length mismatch");
  TORCH_CHECK(g.numel() <= std::numeric_limits<int>::max(), "too many neurons");
}
void index_range(const torch::Tensor& x, int64_t upper, const char* name) {
  if (x.numel()) {
    TORCH_CHECK(x.min().item<int64_t>() >= 0 && x.max().item<int64_t>() < upper,
                name, " index out of bounds");
  }
}

template <typename T> __device__ __forceinline__ T add_rn(T a, T b);
template <> __device__ __forceinline__ float add_rn(float a, float b) { return __fadd_rn(a, b); }
template <> __device__ __forceinline__ double add_rn(double a, double b) { return __dadd_rn(a, b); }
template <typename T> __device__ __forceinline__ T mul_rn(T a, T b);
template <> __device__ __forceinline__ float mul_rn(float a, float b) { return __fmul_rn(a, b); }
template <> __device__ __forceinline__ double mul_rn(double a, double b) { return __dmul_rn(a, b); }

template <typename T>
__global__ void state_kernel(T* v, T* g, int32_t* ref, const uint8_t* requests,
                             uint8_t* spikes, int32_t* counts, int n, T gd, T c0,
                             T cg, T cv, T threshold, T reset, T seed_volt, int ref_steps) {
  const int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= n) return;
  int r = ref[i];
  r = r > 0 ? r - 1 : 0;
  const bool available = r == 0;
  T vi = v[i];
  if (available) {
    const T oldg = g[i];
    const T oldv = vi;
    // Match NumPy float-array expression evaluation: each operation rounds in T.
    g[i] = mul_rn(gd, oldg);
    const T driven = mul_rn(cg, oldg);
    const T leaked = mul_rn(cv, oldv);
    vi = add_rn(c0, add_rn(driven, leaked));
    if (requests[i]) vi = seed_volt;
  }
  const bool fired = available && vi > threshold;
  if (fired) {
    vi = reset;
    g[i] = T(0);
    r = ref_steps;
    counts[i] += 1;
  }
  // Nonavailable neurons retain V/g. Resetting fired cells before arrivals is
  // equivalent to reset-after-arrivals because all their arrivals are rejected.
  v[i] = vi;
  ref[i] = r;
  spikes[i] = fired ? 1 : 0;
}

__global__ void fill_kernel(const int64_t* indptr, const int64_t* sources,
                            const int64_t* offsets, int64_t* output, int n) {
  const int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= n) return;
  const int64_t source = sources[i];
  const int64_t begin = indptr[source];
  const int64_t end = indptr[source + 1];
  for (int64_t j = begin; j < end; ++j) output[offsets[i] + j - begin] = j;
}

template <typename T>
__global__ void arrival_kernel(const int64_t* edges, const int64_t* starts,
                               const int32_t* targets, const T* weights, T* g,
                               const int32_t* ref, int segments, int64_t edge_count) {
  const int segment = blockIdx.x * blockDim.x + threadIdx.x;
  if (segment >= segments) return;
  const int target = targets[segment];
  if (ref[target] != 0) return;
  const int64_t end = segment + 1 < segments ? starts[segment + 1] : edge_count;
  // Start at existing g, never sum weights separately: cancellation can differ.
  T accumulator = g[target];
  for (int64_t j = starts[segment]; j < end; ++j)
    accumulator = add_rn(accumulator, weights[edges[j]]);
  g[target] = accumulator;
}

template <typename T>
__global__ void external_kernel(const T* external, T* g, const int32_t* ref, int n) {
  const int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i < n && ref[i] == 0 && external[i] != T(0))
    g[i] = add_rn(g[i], external[i]);
}
}  // namespace

void author_state_cuda(torch::Tensor v, torch::Tensor g, torch::Tensor ref,
                       torch::Tensor requests, torch::Tensor spikes, torch::Tensor counts,
                       double gd, double c0, double cg, double cv, double threshold,
                       double reset, double seed_volt, int64_t refractory_steps) {
  check_state(g, ref);
  const c10::cuda::CUDAGuard guard(g.device());
  check_tensor(v, g, g.scalar_type(), "v");
  check_tensor(requests, g, torch::kUInt8, "requests");
  check_tensor(spikes, g, torch::kUInt8, "spikes");
  check_tensor(counts, g, torch::kInt32, "counts");
  const int64_t n = g.numel();
  TORCH_CHECK(v.numel() == n && requests.numel() == n && spikes.numel() == n && counts.numel() == n,
              "state vector length mismatch");
  TORCH_CHECK(refractory_steps > 0 && refractory_steps <= std::numeric_limits<int>::max(),
              "invalid refractory length");
  for (double x : {gd, c0, cg, cv, threshold, reset, seed_volt})
    TORCH_CHECK(std::isfinite(x), "nonfinite coefficient");
  if (!n) return;
  const auto stream = at::cuda::getCurrentCUDAStream(g.get_device());
  AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "author_state", [&] {
    state_kernel<scalar_t><<<(n + BLOCK - 1) / BLOCK, BLOCK, 0, stream>>>(
      v.data_ptr<scalar_t>(), g.data_ptr<scalar_t>(), ref.data_ptr<int32_t>(),
      requests.data_ptr<uint8_t>(), spikes.data_ptr<uint8_t>(), counts.data_ptr<int32_t>(),
      static_cast<int>(n), static_cast<scalar_t>(gd), static_cast<scalar_t>(c0),
      static_cast<scalar_t>(cg), static_cast<scalar_t>(cv), static_cast<scalar_t>(threshold),
      static_cast<scalar_t>(reset), static_cast<scalar_t>(seed_volt), static_cast<int>(refractory_steps));
  });
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void author_fill_edges_cuda(torch::Tensor indptr, torch::Tensor sources,
                            torch::Tensor offsets, torch::Tensor output) {
  check_tensor(indptr, indptr, torch::kInt64, "indptr");
  const c10::cuda::CUDAGuard guard(indptr.device());
  check_tensor(sources, indptr, torch::kInt64, "sources");
  check_tensor(offsets, indptr, torch::kInt64, "offsets");
  check_tensor(output, indptr, torch::kInt64, "output");
  TORCH_CHECK(indptr.numel() >= 1 && sources.numel() == offsets.numel(), "CSR shape mismatch");
  TORCH_CHECK(sources.numel() <= std::numeric_limits<int>::max(), "too many sources");
  TORCH_CHECK(indptr[0].item<int64_t>() == 0, "CSR indptr must begin at zero");
  if (indptr.numel() > 1)
    TORCH_CHECK((indptr.slice(0, 1) >= indptr.slice(0, 0, -1)).all().item<bool>(), "CSR indptr not monotonic");
  index_range(sources, indptr.numel() - 1, "source");
  if (!sources.numel()) {
    TORCH_CHECK(output.numel() == 0, "nonempty expansion for no active sources");
    return;
  }
  const auto degrees = indptr.index_select(0, sources + 1) - indptr.index_select(0, sources);
  const auto expected_offsets = degrees.cumsum(0) - degrees;
  TORCH_CHECK(torch::equal(offsets, expected_offsets), "incorrect output offsets");
  TORCH_CHECK(degrees.sum().item<int64_t>() == output.numel(), "incorrect expansion length");
  const auto stream = at::cuda::getCurrentCUDAStream(indptr.get_device());
  fill_kernel<<<(sources.numel() + BLOCK - 1) / BLOCK, BLOCK, 0, stream>>>(
    indptr.data_ptr<int64_t>(), sources.data_ptr<int64_t>(), offsets.data_ptr<int64_t>(),
    output.data_ptr<int64_t>(), static_cast<int>(sources.numel()));
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void author_arrivals_cuda(torch::Tensor edges, torch::Tensor starts, torch::Tensor targets,
                          torch::Tensor weights, torch::Tensor g, torch::Tensor ref) {
  check_state(g, ref);
  const c10::cuda::CUDAGuard guard(g.device());
  check_tensor(edges, g, torch::kInt64, "edges");
  check_tensor(starts, g, torch::kInt64, "starts");
  check_tensor(targets, g, torch::kInt32, "targets");
  check_tensor(weights, g, g.scalar_type(), "weights");
  TORCH_CHECK(starts.numel() == targets.numel(), "segment count mismatch");
  TORCH_CHECK(starts.numel() <= std::numeric_limits<int>::max(), "too many target segments");
  if (!starts.numel()) {
    TORCH_CHECK(!edges.numel(), "edges without target segments");
    return;
  }
  index_range(edges, weights.numel(), "edge");
  index_range(starts, edges.numel(), "segment start");
  index_range(targets, g.numel(), "target");
  TORCH_CHECK(starts[0].item<int64_t>() == 0, "first segment must start at zero");
  if (starts.numel() > 1) {
    TORCH_CHECK((starts.slice(0, 1) > starts.slice(0, 0, -1)).all().item<bool>(), "segment starts not increasing");
    TORCH_CHECK((targets.slice(0, 1) > targets.slice(0, 0, -1)).all().item<bool>(), "targets must be unique and increasing");
  }
  const auto stream = at::cuda::getCurrentCUDAStream(g.get_device());
  AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "author_arrivals", [&] {
    arrival_kernel<scalar_t><<<(starts.numel() + BLOCK - 1) / BLOCK, BLOCK, 0, stream>>>(
      edges.data_ptr<int64_t>(), starts.data_ptr<int64_t>(), targets.data_ptr<int32_t>(),
      weights.data_ptr<scalar_t>(), g.data_ptr<scalar_t>(), ref.data_ptr<int32_t>(),
      static_cast<int>(starts.numel()), edges.numel());
  });
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}

void author_external_cuda(torch::Tensor external, torch::Tensor g, torch::Tensor ref) {
  check_state(g, ref);
  const c10::cuda::CUDAGuard guard(g.device());
  check_tensor(external, g, g.scalar_type(), "external");
  TORCH_CHECK(external.numel() == g.numel(), "external/state length mismatch");
  if (!g.numel()) return;
  const auto stream = at::cuda::getCurrentCUDAStream(g.get_device());
  AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "author_external", [&] {
    external_kernel<scalar_t><<<(g.numel() + BLOCK - 1) / BLOCK, BLOCK, 0, stream>>>(
      external.data_ptr<scalar_t>(), g.data_ptr<scalar_t>(), ref.data_ptr<int32_t>(), static_cast<int>(g.numel()));
  });
  C10_CUDA_KERNEL_LAUNCH_CHECK();
}