#include "akima.hpp"
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <cmath>
#include <stdexcept>

namespace cudakima {

// ====================================================================================
// CUDA Error Checking
// ====================================================================================

#define CUDA_CHECK(call) \
    do { \
        cudaError_t err = call; \
        if (err != cudaSuccess) { \
            throw std::runtime_error(std::string("CUDA error: ") + cudaGetErrorString(err)); \
        } \
    } while(0)

// ====================================================================================
// Check CUDA availability
// ====================================================================================

bool cuda_available() {
    int device_count = 0;
    cudaError_t err = cudaGetDeviceCount(&device_count);
    return (err == cudaSuccess && device_count > 0);
}

// ====================================================================================
// Device Functions - Float
// ====================================================================================

__device__ __forceinline__ float linear_slope_device(const float* x, const float* y, int idx) {
    float dx = x[idx + 1] - x[idx];
    float dy = y[idx + 1] - y[idx];
    return dy / dx;
}

__device__ __forceinline__ float spline_slope_device(
    const float* x, const float* y, const float* linear_slopes,
    int idx, int start, int stop
) {
    // Boundary conditions
    if (idx == start) {
        return (3.0f * linear_slopes[idx] - linear_slopes[idx + 1]) / 2.0f;
    }
    else if (idx == start + 1) {
        float w1 = fabsf(linear_slopes[idx + 1] - linear_slopes[idx]);
        float w2 = fabsf(linear_slopes[idx] - linear_slopes[idx - 1]);
        if (w1 + w2 == 0.0f) {
            return (linear_slopes[idx - 1] + linear_slopes[idx]) / 2.0f;
        }
        return (w1 * linear_slopes[idx - 1] + w2 * linear_slopes[idx]) / (w1 + w2);
    }
    else if (idx == stop - 2) {
        float w1 = fabsf(linear_slopes[idx - 2] - linear_slopes[idx - 1]);
        float w2 = fabsf(linear_slopes[idx - 1] - linear_slopes[idx]);
        if (w1 + w2 == 0.0f) {
            return (linear_slopes[idx - 1] + linear_slopes[idx]) / 2.0f;
        }
        return (w1 * linear_slopes[idx] + w2 * linear_slopes[idx - 1]) / (w1 + w2);
    }
    else if (idx == stop - 1) {
        return (3.0f * linear_slopes[idx - 1] - linear_slopes[idx - 2]) / 2.0f;
    }
    
    // General case
    float mi_2 = linear_slopes[idx - 2];
    float mi_1 = linear_slopes[idx - 1];
    float mi = linear_slopes[idx];
    float mi1 = linear_slopes[idx + 1];
    
    float w1 = fabsf(mi1 - mi);
    float w2 = fabsf(mi_1 - mi_2);
    
    float thr = fmaxf(fmaxf(w1, fabsf(mi - mi_1)), fmaxf(w2, 1e-30f)) * 1e-9f;
    
    if ((w1 + w2) < thr) {
        return (mi_1 + mi) / 2.0f;
    }
    return (w1 * mi_1 + w2 * mi) / (w1 + w2);
}

__device__ __forceinline__ int binary_search_device(const float* x, float target, int start, int stop) {
    int left = start;
    int right = stop - 1;
    
    while (left <= right) {
        int mid = (left + right) / 2;
        if (mid >= stop - 1) {
            return stop - 2;
        }
        if (x[mid] <= target && target < x[mid + 1]) {
            return mid;
        }
        else if (target < x[mid]) {
            right = mid - 1;
        }
        else {
            left = mid + 1;
        }
    }
    return stop - 2;
}

// ====================================================================================
// Device Functions - Double
// ====================================================================================

__device__ __forceinline__ double linear_slope_device_d(const double* x, const double* y, int idx) {
    double dx = x[idx + 1] - x[idx];
    double dy = y[idx + 1] - y[idx];
    return dy / dx;
}

__device__ __forceinline__ double spline_slope_device_d(
    const double* x, const double* y, const double* linear_slopes,
    int idx, int start, int stop
) {
    // Boundary conditions
    if (idx == start) {
        return (3.0 * linear_slopes[idx] - linear_slopes[idx + 1]) / 2.0;
    }
    else if (idx == start + 1) {
        double w1 = fabs(linear_slopes[idx + 1] - linear_slopes[idx]);
        double w2 = fabs(linear_slopes[idx] - linear_slopes[idx - 1]);
        if (w1 + w2 == 0.0) {
            return (linear_slopes[idx - 1] + linear_slopes[idx]) / 2.0;
        }
        return (w1 * linear_slopes[idx - 1] + w2 * linear_slopes[idx]) / (w1 + w2);
    }
    else if (idx == stop - 2) {
        double w1 = fabs(linear_slopes[idx - 2] - linear_slopes[idx - 1]);
        double w2 = fabs(linear_slopes[idx - 1] - linear_slopes[idx]);
        if (w1 + w2 == 0.0) {
            return (linear_slopes[idx - 1] + linear_slopes[idx]) / 2.0;
        }
        return (w1 * linear_slopes[idx] + w2 * linear_slopes[idx - 1]) / (w1 + w2);
    }
    else if (idx == stop - 1) {
        return (3.0 * linear_slopes[idx - 1] - linear_slopes[idx - 2]) / 2.0;
    }
    
    // General case
    double mi_2 = linear_slopes[idx - 2];
    double mi_1 = linear_slopes[idx - 1];
    double mi = linear_slopes[idx];
    double mi1 = linear_slopes[idx + 1];
    
    double w1 = fabs(mi1 - mi);
    double w2 = fabs(mi_1 - mi_2);
    
    double thr = fmax(fmax(w1, fabs(mi - mi_1)), fmax(w2, 1e-99)) * 1e-9;
    
    if ((w1 + w2) < thr) {
        return (mi_1 + mi) / 2.0;
    }
    return (w1 * mi_1 + w2 * mi) / (w1 + w2);
}

__device__ __forceinline__ int binary_search_device_d(const double* x, double target, int start, int stop) {
    int left = start;
    int right = stop - 1;
    
    while (left <= right) {
        int mid = (left + right) / 2;
        if (mid >= stop - 1) {
            return stop - 2;
        }
        if (x[mid] <= target && target < x[mid + 1]) {
            return mid;
        }
        else if (target < x[mid]) {
            right = mid - 1;
        }
        else {
            left = mid + 1;
        }
    }
    return stop - 2;
}

// ====================================================================================
// CUDA Kernels - Float
// ====================================================================================

__global__ void precompute_linear_slopes_kernel(
    const float* x, const float* y, float* linear_slopes,
    int nin, int ngroups, const int* nnans
) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    int group_idx = blockIdx.y;
    
    if (group_idx >= nin) return;
    
    int start = group_idx * ngroups;
    int stop = (group_idx + 1) * ngroups - nnans[group_idx];
    
    if (idx < (stop - start - 1)) {
        int actual_idx = start + idx;
        if (actual_idx < stop - 1) {
            float dx = x[actual_idx + 1] - x[actual_idx];
            float dy = y[actual_idx + 1] - y[actual_idx];
            linear_slopes[actual_idx] = dy / dx;
        }
    }
}

__global__ void precompute_spline_slopes_kernel(
    const float* x, const float* y,
    const float* linear_slopes, float* spline_slopes,
    int nin, int ngroups, const int* nnans
) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    int group_idx = blockIdx.y;
    
    if (group_idx >= nin) return;
    
    int start = group_idx * ngroups;
    int stop = (group_idx + 1) * ngroups - nnans[group_idx];
    
    if (idx < (stop - start) && (stop - start) >= 4) {
        int actual_idx = start + idx;
        spline_slopes[actual_idx] = spline_slope_device(
            x, y, linear_slopes, actual_idx, start, stop
        );
    }
}

__global__ void akima_interpolate_kernel(
    const float* x_new, int n_interp,
    const float* x, const float* y,
    const float* linear_slopes, const float* spline_slopes,
    int nin, int ngroups, const int* nnans,
    float* result, bool cubic
) {
    int i = threadIdx.x + blockIdx.x * blockDim.x;  // x_new index
    int j = blockIdx.y;  // group index
    
    if (i >= n_interp || j >= nin) return;
    
    int start = j * ngroups;
    int stop = (j + 1) * ngroups - nnans[j];
    int special_index = j * n_interp + i;
    
    // Handle edge case
    if (x_new[i] >= x[stop - 1]) {
        result[special_index] = y[stop - 1];
        return;
    }
    
    // Binary search for interval
    int idx = binary_search_device(x, x_new[i], start, stop);
    
    if (!cubic || (stop - start < 4)) {
        // Linear interpolation
        float mi = linear_slopes[idx];
        result[special_index] = y[idx] + mi * (x_new[i] - x[idx]);
    }
    else {
        // Akima spline interpolation
        float mi = linear_slopes[idx];
        float si = spline_slopes[idx];
        float si1 = spline_slopes[idx + 1];
        
        float dx = x[idx + 1] - x[idx];
        float dt = x_new[i] - x[idx];
        
        float p0 = y[idx];
        float p1 = si;
        float p2 = (3.0f * mi - 2.0f * si - si1) / dx;
        float p3 = (si + si1 - 2.0f * mi) / (dx * dx);
        
        result[special_index] = p0 + p1 * dt + p2 * dt * dt + p3 * dt * dt * dt;
    }
}

__global__ void akima_interpolate_multidim_kernel(
    const float* x_new, int n_interp,
    const float* x, const float* y,
    const float* linear_slopes, const float* spline_slopes,
    int nin, int ngroups, const int* nnans,
    float* result, bool cubic
) {
    int i = threadIdx.x + blockIdx.x * blockDim.x;  // x_new index within group
    int j = blockIdx.y;  // group index
    
    if (i >= n_interp || j >= nin) return;
    
    int start = j * ngroups;
    int stop = (j + 1) * ngroups - nnans[j];
    int special_index = j * n_interp + i;
    
    // x_new is multidimensional
    float x_target = x_new[j * n_interp + i];
    
    // Handle edge case
    if (x_target >= x[stop - 1]) {
        result[special_index] = y[stop - 1];
        return;
    }
    
    // Binary search for interval
    int idx = binary_search_device(x, x_target, start, stop);
    
    if (!cubic || (stop - start < 4)) {
        // Linear interpolation
        float mi = linear_slopes[idx];
        result[special_index] = y[idx] + mi * (x_target - x[idx]);
    }
    else {
        // Akima spline interpolation
        float mi = linear_slopes[idx];
        float si = spline_slopes[idx];
        float si1 = spline_slopes[idx + 1];
        
        float dx = x[idx + 1] - x[idx];
        float dt = x_target - x[idx];
        
        float p0 = y[idx];
        float p1 = si;
        float p2 = (3.0f * mi - 2.0f * si - si1) / dx;
        float p3 = (si + si1 - 2.0f * mi) / (dx * dx);
        
        result[special_index] = p0 + p1 * dt + p2 * dt * dt + p3 * dt * dt * dt;
    }
}

// ====================================================================================
// CUDA Kernels - Double
// ====================================================================================

__global__ void precompute_linear_slopes_kernel_d(
    const double* x, const double* y, double* linear_slopes,
    int nin, int ngroups, const int* nnans
) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    int group_idx = blockIdx.y;
    
    if (group_idx >= nin) return;
    
    int start = group_idx * ngroups;
    int stop = (group_idx + 1) * ngroups - nnans[group_idx];
    
    if (idx < (stop - start - 1)) {
        int actual_idx = start + idx;
        if (actual_idx < stop - 1) {
            double dx = x[actual_idx + 1] - x[actual_idx];
            double dy = y[actual_idx + 1] - y[actual_idx];
            linear_slopes[actual_idx] = dy / dx;
        }
    }
}

__global__ void precompute_spline_slopes_kernel_d(
    const double* x, const double* y,
    const double* linear_slopes, double* spline_slopes,
    int nin, int ngroups, const int* nnans
) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    int group_idx = blockIdx.y;
    
    if (group_idx >= nin) return;
    
    int start = group_idx * ngroups;
    int stop = (group_idx + 1) * ngroups - nnans[group_idx];
    
    if (idx < (stop - start) && (stop - start) >= 4) {
        int actual_idx = start + idx;
        spline_slopes[actual_idx] = spline_slope_device_d(
            x, y, linear_slopes, actual_idx, start, stop
        );
    }
}

__global__ void akima_interpolate_kernel_d(
    const double* x_new, int n_interp,
    const double* x, const double* y,
    const double* linear_slopes, const double* spline_slopes,
    int nin, int ngroups, const int* nnans,
    double* result, bool cubic
) {
    int i = threadIdx.x + blockIdx.x * blockDim.x;
    int j = blockIdx.y;
    
    if (i >= n_interp || j >= nin) return;
    
    int start = j * ngroups;
    int stop = (j + 1) * ngroups - nnans[j];
    int special_index = j * n_interp + i;
    
    if (x_new[i] >= x[stop - 1]) {
        result[special_index] = y[stop - 1];
        return;
    }
    
    int idx = binary_search_device_d(x, x_new[i], start, stop);
    
    if (!cubic || (stop - start < 4)) {
        double mi = linear_slopes[idx];
        result[special_index] = y[idx] + mi * (x_new[i] - x[idx]);
    }
    else {
        double mi = linear_slopes[idx];
        double si = spline_slopes[idx];
        double si1 = spline_slopes[idx + 1];
        
        double dx = x[idx + 1] - x[idx];
        double dt = x_new[i] - x[idx];
        
        double p0 = y[idx];
        double p1 = si;
        double p2 = (3.0 * mi - 2.0 * si - si1) / dx;
        double p3 = (si + si1 - 2.0 * mi) / (dx * dx);
        
        result[special_index] = p0 + p1 * dt + p2 * dt * dt + p3 * dt * dt * dt;
    }
}

__global__ void akima_interpolate_multidim_kernel_d(
    const double* x_new, int n_interp,
    const double* x, const double* y,
    const double* linear_slopes, const double* spline_slopes,
    int nin, int ngroups, const int* nnans,
    double* result, bool cubic
) {
    int i = threadIdx.x + blockIdx.x * blockDim.x;
    int j = blockIdx.y;
    
    if (i >= n_interp || j >= nin) return;
    
    int start = j * ngroups;
    int stop = (j + 1) * ngroups - nnans[j];
    int special_index = j * n_interp + i;
    
    double x_target = x_new[j * n_interp + i];
    
    if (x_target >= x[stop - 1]) {
        result[special_index] = y[stop - 1];
        return;
    }
    
    int idx = binary_search_device_d(x, x_target, start, stop);
    
    if (!cubic || (stop - start < 4)) {
        double mi = linear_slopes[idx];
        result[special_index] = y[idx] + mi * (x_target - x[idx]);
    }
    else {
        double mi = linear_slopes[idx];
        double si = spline_slopes[idx];
        double si1 = spline_slopes[idx + 1];
        
        double dx = x[idx + 1] - x[idx];
        double dt = x_target - x[idx];
        
        double p0 = y[idx];
        double p1 = si;
        double p2 = (3.0 * mi - 2.0 * si - si1) / dx;
        double p3 = (si + si1 - 2.0 * mi) / (dx * dx);
        
        result[special_index] = p0 + p1 * dt + p2 * dt * dt + p3 * dt * dt * dt;
    }
}

// ====================================================================================
// Host Functions - Float
// ====================================================================================

void precompute_linear_slopes_cuda(
    const float* x, const float* y, float* linear_slopes,
    int nin, int ngroups, const int* nnans
) {
    float *d_x, *d_y, *d_slopes;
    int *d_nnans;
    
    size_t size_xy = nin * ngroups * sizeof(float);
    size_t size_nnans = nin * sizeof(int);
    
    CUDA_CHECK(cudaMalloc(&d_x, size_xy));
    CUDA_CHECK(cudaMalloc(&d_y, size_xy));
    CUDA_CHECK(cudaMalloc(&d_slopes, size_xy));
    CUDA_CHECK(cudaMalloc(&d_nnans, size_nnans));
    
    CUDA_CHECK(cudaMemcpy(d_x, x, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_y, y, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_nnans, nnans, size_nnans, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemset(d_slopes, 0, size_xy));
    
    int threads = 256;
    int blocks_x = (ngroups + threads - 1) / threads;
    dim3 grid(blocks_x, nin, 1);
    
    precompute_linear_slopes_kernel<<<grid, threads>>>(
        d_x, d_y, d_slopes, nin, ngroups, d_nnans
    );
    CUDA_CHECK(cudaGetLastError());
    
    CUDA_CHECK(cudaMemcpy(linear_slopes, d_slopes, size_xy, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_x));
    CUDA_CHECK(cudaFree(d_y));
    CUDA_CHECK(cudaFree(d_slopes));
    CUDA_CHECK(cudaFree(d_nnans));
}

void precompute_spline_slopes_cuda(
    const float* x, const float* y,
    const float* linear_slopes, float* spline_slopes,
    int nin, int ngroups, const int* nnans
) {
    float *d_x, *d_y, *d_linear, *d_spline;
    int *d_nnans;
    
    size_t size_xy = nin * ngroups * sizeof(float);
    size_t size_nnans = nin * sizeof(int);
    
    CUDA_CHECK(cudaMalloc(&d_x, size_xy));
    CUDA_CHECK(cudaMalloc(&d_y, size_xy));
    CUDA_CHECK(cudaMalloc(&d_linear, size_xy));
    CUDA_CHECK(cudaMalloc(&d_spline, size_xy));
    CUDA_CHECK(cudaMalloc(&d_nnans, size_nnans));
    
    CUDA_CHECK(cudaMemcpy(d_x, x, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_y, y, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_linear, linear_slopes, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_nnans, nnans, size_nnans, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemset(d_spline, 0, size_xy));
    
    int threads = 256;
    int blocks_x = (ngroups + threads - 1) / threads;
    dim3 grid(blocks_x, nin, 1);
    
    precompute_spline_slopes_kernel<<<grid, threads>>>(
        d_x, d_y, d_linear, d_spline, nin, ngroups, d_nnans
    );
    CUDA_CHECK(cudaGetLastError());
    
    CUDA_CHECK(cudaMemcpy(spline_slopes, d_spline, size_xy, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_x));
    CUDA_CHECK(cudaFree(d_y));
    CUDA_CHECK(cudaFree(d_linear));
    CUDA_CHECK(cudaFree(d_spline));
    CUDA_CHECK(cudaFree(d_nnans));
}

void akima_interpolate_cuda(
    const float* x_new, int n_interp,
    const float* x, const float* y,
    const float* linear_slopes, const float* spline_slopes,
    int nin, int ngroups, const int* nnans,
    float* result, bool cubic
) {
    float *d_x_new, *d_x, *d_y, *d_linear, *d_spline, *d_result;
    int *d_nnans;
    
    size_t size_xy = nin * ngroups * sizeof(float);
    size_t size_x_new = n_interp * sizeof(float);
    size_t size_result = nin * n_interp * sizeof(float);
    size_t size_nnans = nin * sizeof(int);
    
    CUDA_CHECK(cudaMalloc(&d_x_new, size_x_new));
    CUDA_CHECK(cudaMalloc(&d_x, size_xy));
    CUDA_CHECK(cudaMalloc(&d_y, size_xy));
    CUDA_CHECK(cudaMalloc(&d_linear, size_xy));
    CUDA_CHECK(cudaMalloc(&d_spline, size_xy));
    CUDA_CHECK(cudaMalloc(&d_result, size_result));
    CUDA_CHECK(cudaMalloc(&d_nnans, size_nnans));
    
    CUDA_CHECK(cudaMemcpy(d_x_new, x_new, size_x_new, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_x, x, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_y, y, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_linear, linear_slopes, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_spline, spline_slopes, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_nnans, nnans, size_nnans, cudaMemcpyHostToDevice));
    
    int threads = 256;
    int blocks_x = (n_interp + threads - 1) / threads;
    dim3 grid(blocks_x, nin, 1);
    
    akima_interpolate_kernel<<<grid, threads>>>(
        d_x_new, n_interp, d_x, d_y, d_linear, d_spline,
        nin, ngroups, d_nnans, d_result, cubic
    );
    CUDA_CHECK(cudaGetLastError());
    
    CUDA_CHECK(cudaMemcpy(result, d_result, size_result, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_x_new));
    CUDA_CHECK(cudaFree(d_x));
    CUDA_CHECK(cudaFree(d_y));
    CUDA_CHECK(cudaFree(d_linear));
    CUDA_CHECK(cudaFree(d_spline));
    CUDA_CHECK(cudaFree(d_result));
    CUDA_CHECK(cudaFree(d_nnans));
}

void akima_interpolate_multidim_cuda(
    const float* x_new, int n_interp,
    const float* x, const float* y,
    const float* linear_slopes, const float* spline_slopes,
    int nin, int ngroups, const int* nnans,
    float* result, bool cubic
) {
    float *d_x_new, *d_x, *d_y, *d_linear, *d_spline, *d_result;
    int *d_nnans;
    
    size_t size_xy = nin * ngroups * sizeof(float);
    size_t size_x_new = nin * n_interp * sizeof(float);
    size_t size_result = nin * n_interp * sizeof(float);
    size_t size_nnans = nin * sizeof(int);
    
    CUDA_CHECK(cudaMalloc(&d_x_new, size_x_new));
    CUDA_CHECK(cudaMalloc(&d_x, size_xy));
    CUDA_CHECK(cudaMalloc(&d_y, size_xy));
    CUDA_CHECK(cudaMalloc(&d_linear, size_xy));
    CUDA_CHECK(cudaMalloc(&d_spline, size_xy));
    CUDA_CHECK(cudaMalloc(&d_result, size_result));
    CUDA_CHECK(cudaMalloc(&d_nnans, size_nnans));
    
    CUDA_CHECK(cudaMemcpy(d_x_new, x_new, size_x_new, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_x, x, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_y, y, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_linear, linear_slopes, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_spline, spline_slopes, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_nnans, nnans, size_nnans, cudaMemcpyHostToDevice));
    
    int threads = 256;
    int blocks_x = (n_interp + threads - 1) / threads;
    dim3 grid(blocks_x, nin, 1);
    
    akima_interpolate_multidim_kernel<<<grid, threads>>>(
        d_x_new, n_interp, d_x, d_y, d_linear, d_spline,
        nin, ngroups, d_nnans, d_result, cubic
    );
    CUDA_CHECK(cudaGetLastError());
    
    CUDA_CHECK(cudaMemcpy(result, d_result, size_result, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_x_new));
    CUDA_CHECK(cudaFree(d_x));
    CUDA_CHECK(cudaFree(d_y));
    CUDA_CHECK(cudaFree(d_linear));
    CUDA_CHECK(cudaFree(d_spline));
    CUDA_CHECK(cudaFree(d_result));
    CUDA_CHECK(cudaFree(d_nnans));
}

// ====================================================================================
// Host Functions - Double
// ====================================================================================

void precompute_linear_slopes_cuda_d(
    const double* x, const double* y, double* linear_slopes,
    int nin, int ngroups, const int* nnans
) {
    double *d_x, *d_y, *d_slopes;
    int *d_nnans;
    
    size_t size_xy = nin * ngroups * sizeof(double);
    size_t size_nnans = nin * sizeof(int);
    
    CUDA_CHECK(cudaMalloc(&d_x, size_xy));
    CUDA_CHECK(cudaMalloc(&d_y, size_xy));
    CUDA_CHECK(cudaMalloc(&d_slopes, size_xy));
    CUDA_CHECK(cudaMalloc(&d_nnans, size_nnans));
    
    CUDA_CHECK(cudaMemcpy(d_x, x, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_y, y, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_nnans, nnans, size_nnans, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemset(d_slopes, 0, size_xy));
    
    int threads = 256;
    int blocks_x = (ngroups + threads - 1) / threads;
    dim3 grid(blocks_x, nin, 1);
    
    precompute_linear_slopes_kernel_d<<<grid, threads>>>(
        d_x, d_y, d_slopes, nin, ngroups, d_nnans
    );
    CUDA_CHECK(cudaGetLastError());
    
    CUDA_CHECK(cudaMemcpy(linear_slopes, d_slopes, size_xy, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_x));
    CUDA_CHECK(cudaFree(d_y));
    CUDA_CHECK(cudaFree(d_slopes));
    CUDA_CHECK(cudaFree(d_nnans));
}

void precompute_spline_slopes_cuda_d(
    const double* x, const double* y,
    const double* linear_slopes, double* spline_slopes,
    int nin, int ngroups, const int* nnans
) {
    double *d_x, *d_y, *d_linear, *d_spline;
    int *d_nnans;
    
    size_t size_xy = nin * ngroups * sizeof(double);
    size_t size_nnans = nin * sizeof(int);
    
    CUDA_CHECK(cudaMalloc(&d_x, size_xy));
    CUDA_CHECK(cudaMalloc(&d_y, size_xy));
    CUDA_CHECK(cudaMalloc(&d_linear, size_xy));
    CUDA_CHECK(cudaMalloc(&d_spline, size_xy));
    CUDA_CHECK(cudaMalloc(&d_nnans, size_nnans));
    
    CUDA_CHECK(cudaMemcpy(d_x, x, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_y, y, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_linear, linear_slopes, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_nnans, nnans, size_nnans, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemset(d_spline, 0, size_xy));
    
    int threads = 256;
    int blocks_x = (ngroups + threads - 1) / threads;
    dim3 grid(blocks_x, nin, 1);
    
    precompute_spline_slopes_kernel_d<<<grid, threads>>>(
        d_x, d_y, d_linear, d_spline, nin, ngroups, d_nnans
    );
    CUDA_CHECK(cudaGetLastError());
    
    CUDA_CHECK(cudaMemcpy(spline_slopes, d_spline, size_xy, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_x));
    CUDA_CHECK(cudaFree(d_y));
    CUDA_CHECK(cudaFree(d_linear));
    CUDA_CHECK(cudaFree(d_spline));
    CUDA_CHECK(cudaFree(d_nnans));
}

void akima_interpolate_cuda_d(
    const double* x_new, int n_interp,
    const double* x, const double* y,
    const double* linear_slopes, const double* spline_slopes,
    int nin, int ngroups, const int* nnans,
    double* result, bool cubic
) {
    double *d_x_new, *d_x, *d_y, *d_linear, *d_spline, *d_result;
    int *d_nnans;
    
    size_t size_xy = nin * ngroups * sizeof(double);
    size_t size_x_new = n_interp * sizeof(double);
    size_t size_result = nin * n_interp * sizeof(double);
    size_t size_nnans = nin * sizeof(int);
    
    CUDA_CHECK(cudaMalloc(&d_x_new, size_x_new));
    CUDA_CHECK(cudaMalloc(&d_x, size_xy));
    CUDA_CHECK(cudaMalloc(&d_y, size_xy));
    CUDA_CHECK(cudaMalloc(&d_linear, size_xy));
    CUDA_CHECK(cudaMalloc(&d_spline, size_xy));
    CUDA_CHECK(cudaMalloc(&d_result, size_result));
    CUDA_CHECK(cudaMalloc(&d_nnans, size_nnans));
    
    CUDA_CHECK(cudaMemcpy(d_x_new, x_new, size_x_new, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_x, x, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_y, y, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_linear, linear_slopes, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_spline, spline_slopes, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_nnans, nnans, size_nnans, cudaMemcpyHostToDevice));
    
    int threads = 256;
    int blocks_x = (n_interp + threads - 1) / threads;
    dim3 grid(blocks_x, nin, 1);
    
    akima_interpolate_kernel_d<<<grid, threads>>>(
        d_x_new, n_interp, d_x, d_y, d_linear, d_spline,
        nin, ngroups, d_nnans, d_result, cubic
    );
    CUDA_CHECK(cudaGetLastError());
    
    CUDA_CHECK(cudaMemcpy(result, d_result, size_result, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_x_new));
    CUDA_CHECK(cudaFree(d_x));
    CUDA_CHECK(cudaFree(d_y));
    CUDA_CHECK(cudaFree(d_linear));
    CUDA_CHECK(cudaFree(d_spline));
    CUDA_CHECK(cudaFree(d_result));
    CUDA_CHECK(cudaFree(d_nnans));
}

void akima_interpolate_multidim_cuda_d(
    const double* x_new, int n_interp,
    const double* x, const double* y,
    const double* linear_slopes, const double* spline_slopes,
    int nin, int ngroups, const int* nnans,
    double* result, bool cubic
) {
    double *d_x_new, *d_x, *d_y, *d_linear, *d_spline, *d_result;
    int *d_nnans;
    
    size_t size_xy = nin * ngroups * sizeof(double);
    size_t size_x_new = nin * n_interp * sizeof(double);
    size_t size_result = nin * n_interp * sizeof(double);
    size_t size_nnans = nin * sizeof(int);
    
    CUDA_CHECK(cudaMalloc(&d_x_new, size_x_new));
    CUDA_CHECK(cudaMalloc(&d_x, size_xy));
    CUDA_CHECK(cudaMalloc(&d_y, size_xy));
    CUDA_CHECK(cudaMalloc(&d_linear, size_xy));
    CUDA_CHECK(cudaMalloc(&d_spline, size_xy));
    CUDA_CHECK(cudaMalloc(&d_result, size_result));
    CUDA_CHECK(cudaMalloc(&d_nnans, size_nnans));
    
    CUDA_CHECK(cudaMemcpy(d_x_new, x_new, size_x_new, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_x, x, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_y, y, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_linear, linear_slopes, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_spline, spline_slopes, size_xy, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_nnans, nnans, size_nnans, cudaMemcpyHostToDevice));
    
    int threads = 256;
    int blocks_x = (n_interp + threads - 1) / threads;
    dim3 grid(blocks_x, nin, 1);
    
    akima_interpolate_multidim_kernel_d<<<grid, threads>>>(
        d_x_new, n_interp, d_x, d_y, d_linear, d_spline,
        nin, ngroups, d_nnans, d_result, cubic
    );
    CUDA_CHECK(cudaGetLastError());
    
    CUDA_CHECK(cudaMemcpy(result, d_result, size_result, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_x_new));
    CUDA_CHECK(cudaFree(d_x));
    CUDA_CHECK(cudaFree(d_y));
    CUDA_CHECK(cudaFree(d_linear));
    CUDA_CHECK(cudaFree(d_spline));
    CUDA_CHECK(cudaFree(d_result));
    CUDA_CHECK(cudaFree(d_nnans));
}

// ====================================================================================
// Template Class Implementation
// ====================================================================================

template<typename T>
std::vector<T> AkimaInterpolator<T>::interpolate(
    const std::vector<T>& x_new,
    const std::vector<T>& x_flat,
    const std::vector<T>& y_flat,
    int nin, int ngroups,
    const std::vector<int>& nnans
) {
    int n_interp = static_cast<int>(x_new.size());
    int total_points = nin * ngroups;
    
    std::vector<T> linear_slopes(total_points, 0);
    std::vector<T> spline_slopes(total_points, 0);
    std::vector<T> result(nin * n_interp, 0);
    
    bool cubic = (order_ == Order::CUBIC);
    
    if (use_gpu_) {
        // GPU path
        if constexpr (std::is_same_v<T, float>) {
            precompute_linear_slopes_cuda(
                x_flat.data(), y_flat.data(), linear_slopes.data(),
                nin, ngroups, nnans.data()
            );
            if (cubic) {
                precompute_spline_slopes_cuda(
                    x_flat.data(), y_flat.data(),
                    linear_slopes.data(), spline_slopes.data(),
                    nin, ngroups, nnans.data()
                );
            }
            akima_interpolate_cuda(
                x_new.data(), n_interp,
                x_flat.data(), y_flat.data(),
                linear_slopes.data(), spline_slopes.data(),
                nin, ngroups, nnans.data(),
                result.data(), cubic
            );
        } else {
            precompute_linear_slopes_cuda_d(
                x_flat.data(), y_flat.data(), linear_slopes.data(),
                nin, ngroups, nnans.data()
            );
            if (cubic) {
                precompute_spline_slopes_cuda_d(
                    x_flat.data(), y_flat.data(),
                    linear_slopes.data(), spline_slopes.data(),
                    nin, ngroups, nnans.data()
                );
            }
            akima_interpolate_cuda_d(
                x_new.data(), n_interp,
                x_flat.data(), y_flat.data(),
                linear_slopes.data(), spline_slopes.data(),
                nin, ngroups, nnans.data(),
                result.data(), cubic
            );
        }
    } else {
        // CPU path
        precompute_linear_slopes_cpu(
            x_flat.data(), y_flat.data(), linear_slopes.data(),
            nin, ngroups, nnans.data()
        );
        if (cubic) {
            precompute_spline_slopes_cpu(
                x_flat.data(), y_flat.data(),
                linear_slopes.data(), spline_slopes.data(),
                nin, ngroups, nnans.data()
            );
        }
        akima_interpolate_cpu(
            x_new.data(), n_interp,
            x_flat.data(), y_flat.data(),
            linear_slopes.data(), spline_slopes.data(),
            nin, ngroups, nnans.data(),
            result.data(), cubic
        );
    }
    
    return result;
}

template<typename T>
std::vector<T> AkimaInterpolator<T>::interpolate_multidim(
    const std::vector<T>& x_new,
    const std::vector<T>& x_flat,
    const std::vector<T>& y_flat,
    int nin, int ngroups, int n_interp,
    const std::vector<int>& nnans
) {
    int total_points = nin * ngroups;
    
    std::vector<T> linear_slopes(total_points, 0);
    std::vector<T> spline_slopes(total_points, 0);
    std::vector<T> result(nin * n_interp, 0);
    
    bool cubic = (order_ == Order::CUBIC);
    
    if (use_gpu_) {
        // GPU path
        if constexpr (std::is_same_v<T, float>) {
            precompute_linear_slopes_cuda(
                x_flat.data(), y_flat.data(), linear_slopes.data(),
                nin, ngroups, nnans.data()
            );
            if (cubic) {
                precompute_spline_slopes_cuda(
                    x_flat.data(), y_flat.data(),
                    linear_slopes.data(), spline_slopes.data(),
                    nin, ngroups, nnans.data()
                );
            }
            akima_interpolate_multidim_cuda(
                x_new.data(), n_interp,
                x_flat.data(), y_flat.data(),
                linear_slopes.data(), spline_slopes.data(),
                nin, ngroups, nnans.data(),
                result.data(), cubic
            );
        } else {
            precompute_linear_slopes_cuda_d(
                x_flat.data(), y_flat.data(), linear_slopes.data(),
                nin, ngroups, nnans.data()
            );
            if (cubic) {
                precompute_spline_slopes_cuda_d(
                    x_flat.data(), y_flat.data(),
                    linear_slopes.data(), spline_slopes.data(),
                    nin, ngroups, nnans.data()
                );
            }
            akima_interpolate_multidim_cuda_d(
                x_new.data(), n_interp,
                x_flat.data(), y_flat.data(),
                linear_slopes.data(), spline_slopes.data(),
                nin, ngroups, nnans.data(),
                result.data(), cubic
            );
        }
    } else {
        // CPU path
        precompute_linear_slopes_cpu(
            x_flat.data(), y_flat.data(), linear_slopes.data(),
            nin, ngroups, nnans.data()
        );
        if (cubic) {
            precompute_spline_slopes_cpu(
                x_flat.data(), y_flat.data(),
                linear_slopes.data(), spline_slopes.data(),
                nin, ngroups, nnans.data()
            );
        }
        akima_interpolate_multidim_cpu(
            x_new.data(), n_interp,
            x_flat.data(), y_flat.data(),
            linear_slopes.data(), spline_slopes.data(),
            nin, ngroups, nnans.data(),
            result.data(), cubic
        );
    }
    
    return result;
}

// Explicit template instantiations
template class AkimaInterpolator<float>;
template class AkimaInterpolator<double>;

} // namespace cudakima
