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

namespace {
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");
}

__global__ void state_kernel(float* v, float* g, int* refractory,
                             float* arrival, unsigned char* spikes,
                             int* counts, int n, float gd, float alpha,
                             float rest, float reset, float threshold,
                             int ref_steps) {
  int i=blockIdx.x*blockDim.x+threadIdx.x;
  if(i>=n) return;
  float gi=__fadd_rn(__fmul_rn(g[i],gd),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){
    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]=ref_steps;counts[i]+=1;}
  v[i]=vi;spikes[i]=fired?1:0;
}

__global__ void force_kernel(float* v,int* refractory,unsigned char* spikes,
                             int* counts,const int64_t* targets,
                             const unsigned char* forced,int n,float reset,
                             int ref_steps){
  int j=blockIdx.x*blockDim.x+threadIdx.x;
  if(j>=n||!forced[j]) return;
  int64_t i=targets[j]; if(spikes[i]||refractory[i]!=0) return;
  spikes[i]=1;v[i]=reset;refractory[i]=ref_steps;counts[i]+=1;
}

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

// Exactly one CUDA thread owns a target segment.  The loop order is the
// canonical target -> source -> edge_id order supplied by the host.
__global__ void segmented_sum_kernel(const int64_t* edges,const int64_t* starts,
                                     const int* targets,const float* weights,
                                     float* destination,int segments,
                                     int64_t edge_count){
  int s=blockIdx.x*blockDim.x+threadIdx.x;
  if(s>=segments) return;
  int64_t begin=starts[s],end=(s+1<segments)?starts[s+1]:edge_count;
  float acc=0.0f;
  for(int64_t j=begin;j<end;++j) acc=__fadd_rn(acc,weights[edges[j]]);
  destination[targets[s]]=__fadd_rn(destination[targets[s]],acc);
}
}

void state_step_cuda(torch::Tensor v,torch::Tensor g,torch::Tensor refractory,
                     torch::Tensor arrival,torch::Tensor spikes,torch::Tensor counts,
                     double gd,double alpha,double rest,double reset,double threshold,
                     int64_t ref_steps){
  check_cuda(v);check_cuda(g);check_cuda(refractory);check_cuda(arrival);check_cuda(spikes);check_cuda(counts);
  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>(),counts.data_ptr<int>(),n,(float)gd,(float)alpha,(float)rest,(float)reset,(float)threshold,(int)ref_steps);C10_CUDA_KERNEL_LAUNCH_CHECK();
}
void force_step_cuda(torch::Tensor v,torch::Tensor refractory,torch::Tensor spikes,
                     torch::Tensor counts,torch::Tensor targets,torch::Tensor forced,
                     double reset,int64_t ref_steps){
  check_cuda(targets);check_cuda(forced);int n=targets.numel();force_kernel<<<(n+255)/256,256>>>(v.data_ptr<float>(),refractory.data_ptr<int>(),spikes.data_ptr<unsigned char>(),counts.data_ptr<int>(),targets.data_ptr<int64_t>(),forced.data_ptr<unsigned char>(),n,(float)reset,(int)ref_steps);C10_CUDA_KERNEL_LAUNCH_CHECK();
}
void fill_edges_cuda(torch::Tensor indptr,torch::Tensor sources,torch::Tensor offsets,torch::Tensor output){
  check_cuda(indptr);check_cuda(sources);check_cuda(offsets);check_cuda(output);int n=sources.numel();fill_edges_kernel<<<(n+255)/256,256>>>(indptr.data_ptr<int64_t>(),sources.data_ptr<int64_t>(),offsets.data_ptr<int64_t>(),output.data_ptr<int64_t>(),n);C10_CUDA_KERNEL_LAUNCH_CHECK();
}
void segmented_sum_cuda(torch::Tensor edges,torch::Tensor starts,torch::Tensor targets,
                        torch::Tensor weights,torch::Tensor destination){
  check_cuda(edges);check_cuda(starts);check_cuda(targets);check_cuda(weights);check_cuda(destination);
  TORCH_CHECK(edges.scalar_type()==torch::kInt64&&starts.scalar_type()==torch::kInt64&&targets.scalar_type()==torch::kInt32,"index dtype mismatch");
  int n=targets.numel();if(n) segmented_sum_kernel<<<(n+255)/256,256>>>(edges.data_ptr<int64_t>(),starts.data_ptr<int64_t>(),targets.data_ptr<int>(),weights.data_ptr<float>(),destination.data_ptr<float>(),n,edges.numel());C10_CUDA_KERNEL_LAUNCH_CHECK();
}