// CUDA Driver/NVRTC support for a restricted integer AltAILIF/Delta KNP backend.
#pragma once
#include <cuda.h>
#include <nvrtc.h>
#include <cstdint>
#include <stdexcept>
#include <string>
#include <vector>

namespace knp_altai_cuda {
inline void checked(CUresult code) {
    if (code != CUDA_SUCCESS) {
        const char *message = nullptr;
        cuGetErrorString(code, &message);
        throw std::runtime_error(std::string("AltAI CUDA: ") + (message ? message : "driver error"));
    }
}
struct NeuronState {
    float potential, pre;
    std::int64_t activity;
    int threshold, negative, leak, reset;
    unsigned flags;
};
static_assert(sizeof(NeuronState) == 40, "CUDA neuron ABI");
inline const char *kernel_source = R"CUDA(
struct NeuronState { float potential, pre; long long activity; int threshold, negative, leak, reset; unsigned flags; };
extern "C" __global__ void population(NeuronState *neurons, long long *pending, long long *absolute,
                                     long long *events, unsigned char *spikes, int *error,
                                     int n, int slots, unsigned long long step) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    NeuronState state = neurons[i];
    float v = (state.flags & 16) ? (float)state.reset : state.potential;
    state.pre = v;
    unsigned long long index = (step % slots) * (unsigned long long)n + i;
    long long bound = (long long)(v < 0 ? -v : v) + absolute[index] + (state.leak < 0 ? -state.leak : state.leak);
    if (bound >= 16777216LL) { atomicExch(error, 1); spikes[i] = 0; return; }
    if (events[index]) v += (float)pending[index];
    pending[index] = absolute[index] = events[index] = 0;
    v += state.leak * (((state.flags & 4) && v < 0) ? -1 : 1);
    if (state.activity > 0) --state.activity;
    else if (state.activity < 0) ++state.activity;
    if (state.activity == 0) state.activity = 9223372036854775807LL;
    bool fired = false, reset = false;
    if (v >= state.threshold) {
        fired = state.activity > 0;
        if (state.flags & 1) v -= state.threshold;
        if (state.flags & 2) { v = (float)state.reset; reset = true; }
    }
    if (v <= -(float)state.negative && !reset) {
        if (state.flags & 8) v = -(float)state.negative;
        else if (state.flags & 2) v = (float)(-state.reset);
        else if (state.flags & 1) v += state.negative;
    }
    state.potential = v;
    neurons[i] = state;
    spikes[i] = fired;
}
extern "C" __global__ void projection(const unsigned *rows, const unsigned *sources,
                                     const int *weights, const unsigned *delays, const int *counts,
                                     long long *pending, long long *absolute, long long *events,
                                     int n, int slots, unsigned long long step) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    for (unsigned s = rows[i]; s < rows[i + 1]; ++s) {
        int count = counts[sources[s]];
        if (!count) continue;
        long long impact = (long long)weights[s] * count;
        unsigned long long index = ((step + delays[s]) % slots) * (unsigned long long)n + i;
        pending[index] += impact;
        absolute[index] += impact < 0 ? -impact : impact;
        events[index] += count;
    }
}

extern "C" __global__ void projection_wide(const unsigned *rows, const unsigned *sources,
                                     const int *weights, const unsigned *delays, const int *counts,
                                     long long *pending, long long *absolute, long long *events,
                                     int n, int slots, unsigned long long step) {
    int i = blockIdx.x, t = threadIdx.x;
    if (i >= n) return;
    __shared__ long long sums[256], absums[256], event_sums[256];
    for (int delay=1;delay<slots;++delay) {
        long long value=0, ab=0, ev=0;
        for (unsigned s=rows[i]+t;s<rows[i+1];s+=256) {
            if (delays[s] != (unsigned)delay) continue;
            int count=counts[sources[s]];
            long long impact=(long long)weights[s]*count;
            value+=impact;ab+=impact<0?-impact:impact;ev+=count;
        }
        sums[t]=value;absums[t]=ab;event_sums[t]=ev;
        __syncthreads();
        for (int shift=128;shift;shift/=2) {
            if(t<shift) {sums[t]+=sums[t+shift];absums[t]+=absums[t+shift];event_sums[t]+=event_sums[t+shift];}
            __syncthreads();
        }
        if(t==0) {
            unsigned long long index=((step+delay)%slots)*(unsigned long long)n+i;
            pending[index]+=sums[0];absolute[index]+=absums[0];events[index]+=event_sums[0];
        }
        __syncthreads();
    }
}
)CUDA";

class Runtime {
public:
    CUcontext context = nullptr;
    CUdevice device = 0;
    CUmodule module = nullptr;
    CUfunction population_kernel = nullptr, projection_kernel = nullptr, wide_kernel = nullptr;
    std::vector<CUdeviceptr> allocations;
    Runtime() = default;
    Runtime(const Runtime &) = delete;
    Runtime &operator=(const Runtime &) = delete;
    ~Runtime() {
        if (context) {
            cuCtxSetCurrent(context);
            cuCtxSynchronize();
            for (auto pointer : allocations) cuMemFree(pointer);
            if (module) cuModuleUnload(module);
            cuDevicePrimaryCtxRelease(device);
        }
    }
    void initialize(int index) {
        checked(cuInit(0));
        checked(cuDeviceGet(&device, index));
        checked(cuDevicePrimaryCtxRetain(&context, device));
        current();
        int major = 0, minor = 0;
        checked(cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device));
        checked(cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device));
        std::string architecture = "--gpu-architecture=compute_" + std::to_string(major) + std::to_string(minor);
        nvrtcProgram program = nullptr;
        auto code = nvrtcCreateProgram(&program, kernel_source, "knp_altai.cu", 0, nullptr, nullptr);
        if (code != NVRTC_SUCCESS) throw std::runtime_error(nvrtcGetErrorString(code));
        const char *options[] = {architecture.c_str(), "--std=c++11", "--fmad=false"};
        code = nvrtcCompileProgram(program, 3, options);
        if (code != NVRTC_SUCCESS) {
            size_t length = 0;
            nvrtcGetProgramLogSize(program, &length);
            std::string log(length, '\0');
            nvrtcGetProgramLog(program, log.data());
            nvrtcDestroyProgram(&program);
            throw std::runtime_error("NVRTC: " + log);
        }
        size_t length = 0;
        auto result = nvrtcGetPTXSize(program, &length);
        if (result != NVRTC_SUCCESS) { nvrtcDestroyProgram(&program); throw std::runtime_error(nvrtcGetErrorString(result)); }
        std::vector<char> ptx(length);
        result = nvrtcGetPTX(program, ptx.data());
        nvrtcDestroyProgram(&program);
        if (result != NVRTC_SUCCESS) throw std::runtime_error(nvrtcGetErrorString(result));
        checked(cuModuleLoadData(&module, ptx.data()));
        checked(cuModuleGetFunction(&population_kernel, module, "population"));
        checked(cuModuleGetFunction(&projection_kernel, module, "projection"));
        checked(cuModuleGetFunction(&wide_kernel, module, "projection_wide"));
    }
    void current() const { checked(cuCtxSetCurrent(context)); }
    CUdeviceptr allocate(size_t bytes, const void *data = nullptr) {
        CUdeviceptr result = 0;
        checked(cuMemAlloc(&result, bytes ? bytes : 4));
        allocations.push_back(result);
        if (data && bytes) checked(cuMemcpyHtoD(result, data, bytes));
        else checked(cuMemsetD8(result, 0, bytes ? bytes : 4));
        return result;
    }
    void launch(CUfunction function, int n, void **arguments) {
        if (n) checked(cuLaunchKernel(function, (n + 255) / 256, 1, 1, 256, 1, 1, 0, nullptr, arguments, nullptr));
    }
};
} // namespace knp_altai_cuda