#pragma once

#include <vector>
#include <cstdint>
#include <stdexcept>
#include <cmath>
#include <algorithm>

namespace cudakima {

// Forward declarations for CUDA functions
void precompute_linear_slopes_cuda(
    const float* x, const float* y, float* linear_slopes,
    int nin, int ngroups, const int* 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
);

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
);

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
);

// Double precision versions
void precompute_linear_slopes_cuda_d(
    const double* x, const double* y, double* linear_slopes,
    int nin, int ngroups, const int* 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
);

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
);

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
);

// CPU fallback implementations
template<typename T>
T linear_slope_cpu(const T* x, const T* y, int idx) {
    T dx = x[idx + 1] - x[idx];
    T dy = y[idx + 1] - y[idx];
    return dy / dx;
}

template<typename T>
T spline_slope_cpu(const T* x, const T* y, const T* linear_slopes, 
                   int idx, int start, int stop) {
    // Boundary conditions
    if (idx == start) {
        return (3 * linear_slopes[idx] - linear_slopes[idx + 1]) / 2;
    }
    else if (idx == start + 1) {
        T w1 = std::abs(linear_slopes[idx + 1] - linear_slopes[idx]);
        T w2 = std::abs(linear_slopes[idx] - linear_slopes[idx - 1]);
        if (w1 + w2 == 0) {
            return (linear_slopes[idx - 1] + linear_slopes[idx]) / 2;
        }
        return (w1 * linear_slopes[idx - 1] + w2 * linear_slopes[idx]) / (w1 + w2);
    }
    else if (idx == stop - 2) {
        T w1 = std::abs(linear_slopes[idx - 2] - linear_slopes[idx - 1]);
        T w2 = std::abs(linear_slopes[idx - 1] - linear_slopes[idx]);
        if (w1 + w2 == 0) {
            return (linear_slopes[idx - 1] + linear_slopes[idx]) / 2;
        }
        return (w1 * linear_slopes[idx] + w2 * linear_slopes[idx - 1]) / (w1 + w2);
    }
    else if (idx == stop - 1) {
        return (3 * linear_slopes[idx - 1] - linear_slopes[idx - 2]) / 2;
    }
    
    // General case
    T mi_2 = linear_slopes[idx - 2];
    T mi_1 = linear_slopes[idx - 1];
    T mi = linear_slopes[idx];
    T mi1 = linear_slopes[idx + 1];
    
    T w1 = std::abs(mi1 - mi);
    T w2 = std::abs(mi_1 - mi_2);
    
    T thr = std::max({w1, std::abs(mi - mi_1), w2, static_cast<T>(1e-99)}) * static_cast<T>(1e-9);
    
    if ((w1 + w2) < thr) {
        return (mi_1 + mi) / 2;
    }
    return (w1 * mi_1 + w2 * mi) / (w1 + w2);
}

template<typename T>
int binary_search_cpu(const T* x, T 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;
}

template<typename T>
void precompute_linear_slopes_cpu(
    const T* x, const T* y, T* linear_slopes,
    int nin, int ngroups, const int* nnans
) {
    for (int j = 0; j < nin; ++j) {
        int start = j * ngroups;
        int stop = (j + 1) * ngroups - nnans[j];
        
        for (int i = start; i < stop - 1; ++i) {
            T dx = x[i + 1] - x[i];
            T dy = y[i + 1] - y[i];
            linear_slopes[i] = dy / dx;
        }
    }
}

template<typename T>
void precompute_spline_slopes_cpu(
    const T* x, const T* y, 
    const T* linear_slopes, T* spline_slopes,
    int nin, int ngroups, const int* nnans
) {
    for (int j = 0; j < nin; ++j) {
        int start = j * ngroups;
        int stop = (j + 1) * ngroups - nnans[j];
        
        if (stop - start >= 4) {
            for (int i = start; i < stop; ++i) {
                spline_slopes[i] = spline_slope_cpu(x, y, linear_slopes, i, start, stop);
            }
        }
    }
}

template<typename T>
void akima_interpolate_cpu(
    const T* x_new, int n_interp,
    const T* x, const T* y,
    const T* linear_slopes, const T* spline_slopes,
    int nin, int ngroups, const int* nnans,
    T* result, bool cubic
) {
    #pragma omp parallel for collapse(2)
    for (int j = 0; j < nin; ++j) {
        for (int i = 0; i < n_interp; ++i) {
            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];
                continue;
            }
            
            // Binary search for interval
            int idx = binary_search_cpu(x, x_new[i], start, stop);
            
            if (!cubic || (stop - start < 4)) {
                // Linear interpolation
                T mi = linear_slopes[idx];
                result[special_index] = y[idx] + mi * (x_new[i] - x[idx]);
            }
            else {
                // Akima spline interpolation
                T mi = linear_slopes[idx];
                T si = spline_slopes[idx];
                T si1 = spline_slopes[idx + 1];
                
                T dx = x[idx + 1] - x[idx];
                T dt = x_new[i] - x[idx];
                
                T p0 = y[idx];
                T p1 = si;
                T p2 = (3 * mi - 2 * si - si1) / dx;
                T p3 = (si + si1 - 2 * mi) / (dx * dx);
                
                result[special_index] = p0 + p1 * dt + p2 * dt * dt + p3 * dt * dt * dt;
            }
        }
    }
}

template<typename T>
void akima_interpolate_multidim_cpu(
    const T* x_new, int n_interp,
    const T* x, const T* y,
    const T* linear_slopes, const T* spline_slopes,
    int nin, int ngroups, const int* nnans,
    T* result, bool cubic
) {
    #pragma omp parallel for collapse(2)
    for (int j = 0; j < nin; ++j) {
        for (int i = 0; i < n_interp; ++i) {
            int start = j * ngroups;
            int stop = (j + 1) * ngroups - nnans[j];
            int special_index = j * n_interp + i;
            
            // x_new is multidimensional: x_new[j * n_interp + i]
            T x_target = x_new[j * n_interp + i];
            
            // Handle edge case
            if (x_target >= x[stop - 1]) {
                result[special_index] = y[stop - 1];
                continue;
            }
            
            // Binary search for interval
            int idx = binary_search_cpu(x, x_target, start, stop);
            
            if (!cubic || (stop - start < 4)) {
                // Linear interpolation
                T mi = linear_slopes[idx];
                result[special_index] = y[idx] + mi * (x_target - x[idx]);
            }
            else {
                // Akima spline interpolation
                T mi = linear_slopes[idx];
                T si = spline_slopes[idx];
                T si1 = spline_slopes[idx + 1];
                
                T dx = x[idx + 1] - x[idx];
                T dt = x_target - x[idx];
                
                T p0 = y[idx];
                T p1 = si;
                T p2 = (3 * mi - 2 * si - si1) / dx;
                T p3 = (si + si1 - 2 * mi) / (dx * dx);
                
                result[special_index] = p0 + p1 * dt + p2 * dt * dt + p3 * dt * dt * dt;
            }
        }
    }
}

// Check if CUDA is available at runtime
bool cuda_available();

// Main interpolation class
template<typename T>
class AkimaInterpolator {
public:
    enum class Order { LINEAR, CUBIC };
    
    AkimaInterpolator(bool use_gpu = true, Order order = Order::CUBIC, 
                      int threads_per_block = 256)
        : use_gpu_(use_gpu && cuda_available())
        , order_(order)
        , threads_per_block_(threads_per_block)
    {}
    
    void set_use_gpu(bool use_gpu) { use_gpu_ = use_gpu && cuda_available(); }
    void set_order(Order order) { order_ = order; }
    void set_threads_per_block(int tpb) { threads_per_block_ = tpb; }
    
    bool using_gpu() const { return use_gpu_; }
    Order get_order() const { return order_; }
    
    // Main interpolation method - 1D x_new (same interpolation points for all groups)
    std::vector<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
    );
    
    // Multidimensional x_new (different interpolation points for each group)
    std::vector<T> interpolate_multidim(
        const std::vector<T>& x_new,  // shape: (nin * n_interp)
        const std::vector<T>& x_flat,
        const std::vector<T>& y_flat,
        int nin, int ngroups, int n_interp,
        const std::vector<int>& nnans
    );

private:
    bool use_gpu_;
    Order order_;
    int threads_per_block_;
};

// Explicit template instantiations (defined in .cu file)
extern template class AkimaInterpolator<float>;
extern template class AkimaInterpolator<double>;

} // namespace cudakima
