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

// Generic CUDA kernel for 1D convolution.
// Computes output y of shape [batch, out_channels, length_out]
// where length_out = length_in - kernel_size + 1.
__global__ void conv1d_forward_kernel(const float* __restrict__ x,
                                        const float* __restrict__ w,
                                        float* __restrict__ y,
                                        int batch,
                                        int in_channels,
                                        int length_in,
                                        int out_channels,
                                        int kernel_size,
                                        int length_out) {
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch * out_channels * length_out;
    if (index < total) {
        // Determine indices: p: spatial output index, m: output channel, b: batch index.
        int p = index % length_out;
        int tmp = index / length_out;
        int m = tmp % out_channels;
        int b = tmp / out_channels;
        float sum = 0.0f;
        // Loop over input channels and kernel positions.
        for (int c = 0; c < in_channels; c++) {
            for (int k = 0; k < kernel_size; k++) {
                // Calculate the index into the input tensor.
                int x_index = b * (in_channels * length_in) + c * length_in + (p + k);
                // Calculate the index into the weight tensor.
                int w_index = m * (in_channels * kernel_size) + c * kernel_size + k;
                sum += x[x_index] * w[w_index];
            }
        }
        y[index] = sum;
    }
}

// Optimized unrolled CUDA kernel for 1D convolution when in_channels==3 and kernel_size==3.
__global__ void conv1d_forward_unrolled_kernel(const float* __restrict__ x,
                                                 const float* __restrict__ w,
                                                 float* __restrict__ y,
                                                 int batch,
                                                 int length_in,
                                                 int out_channels,
                                                 int length_out) {
    // Assumes in_channels == 3 and kernel_size == 3.
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch * out_channels * length_out;
    if (index < total) {
        int p = index % length_out;
        int tmp = index / length_out;
        int m = tmp % out_channels;
        int b = tmp / out_channels;
        
        int stride = length_in;  // Each channel has "length_in" elements.
        int base_x = b * (3 * length_in); // for in_channels == 3.
        int base_w = m * (3 * 3);          // for kernel_size == 3 and in_channels == 3.
        
        float sum = 0.0f;
        // Manually unrolled computation for channel 0.
        sum += x[base_x + 0 * stride + (p + 0)] * w[base_w + 0];
        sum += x[base_x + 0 * stride + (p + 1)] * w[base_w + 1];
        sum += x[base_x + 0 * stride + (p + 2)] * w[base_w + 2];
        // Channel 1.
        sum += x[base_x + 1 * stride + (p + 0)] * w[base_w + 3];
        sum += x[base_x + 1 * stride + (p + 1)] * w[base_w + 4];
        sum += x[base_x + 1 * stride + (p + 2)] * w[base_w + 5];
        // Channel 2.
        sum += x[base_x + 2 * stride + (p + 0)] * w[base_w + 6];
        sum += x[base_x + 2 * stride + (p + 1)] * w[base_w + 7];
        sum += x[base_x + 2 * stride + (p + 2)] * w[base_w + 8];
        
        y[index] = sum;
    }
}

torch::Tensor conv1d_forward(torch::Tensor x, torch::Tensor w) {
    // Enforce proper CUDA device using c10::cuda::CUDAGuard.
    c10::cuda::CUDAGuard device_guard(x.device());
    
    TORCH_CHECK(x.dim() == 3, "Input must be a 3D tensor");
    TORCH_CHECK(w.dim() == 3, "Weight must be a 3D tensor");
    
    int batch = x.size(0);
    int in_channels = x.size(1);
    int length_in = x.size(2);
    
    int out_channels = w.size(0);
    int kernel_size = w.size(2);
    
    int length_out = length_in - kernel_size + 1;
    TORCH_CHECK(length_out > 0, "Invalid output length. Check kernel_size and input length.");
    
    auto y = torch::empty({batch, out_channels, length_out}, x.options());
    int total = batch * out_channels * length_out;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    
    // Use specialized unrolled kernel if possible.
    if (in_channels == 3 && kernel_size == 3) {
        conv1d_forward_unrolled_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(),
            w.data_ptr<float>(),
            y.data_ptr<float>(),
            batch,
            length_in,
            out_channels,
            length_out
        );
    } else {
        conv1d_forward_kernel<<<blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
            x.data_ptr<float>(),
            w.data_ptr<float>(),
            y.data_ptr<float>(),
            batch,
            in_channels,
            length_in,
            out_channels,
            kernel_size,
            length_out
        );
    }
    
    return y;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("conv1d_forward", &conv1d_forward, "1D convolution forward (CUDA)");
}