import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
import math

cpp_source = """
#include <torch/extension.h>

torch::Tensor lppool1d_cuda_forward(
    const torch::Tensor& input,
    int norm_power,
    int kernel_size,
    int stride
);
"""

# CUDA 源代码
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>

template <typename T>
__global__ void lppool1d_kernel(
    T* output,
    const T* input,
    const T norm_power,
    const int kernel_size,
    const int stride,
    const int N, const int C, const int L_in, const int L_out,
    const int64_t n_elements_out)
{
    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n_elements_out) return;

    // --- 1D 输出索引映射回 3D 输出坐标 (n, c, l_out) ---
    const int l_out = idx % L_out;
    const int c = (idx / L_out) % C;
    const int n = idx / (L_out * C);
    
    // --- 计算输入窗口的范围 ---
    const int start = l_out * stride;
    const int end = fminf(start + kernel_size, L_in);

    // --- 在寄存器中循环累加 ---
    T sum_val = 0.0f;
    const T* input_channel_ptr = input + (n * C + c) * L_in;

    for (int k = start; k < end; ++k) {
        sum_val += powf(input_channel_ptr[k], norm_power);
    }
    
    // --- 最终计算并写入结果 ---
    output[idx] = powf(sum_val, 1.0f / norm_power);
}

torch::Tensor lppool1d_cuda_forward(
    const torch::Tensor& input,
    const int norm_power,
    const int kernel_size,
    const int stride)
{
    TORCH_CHECK(input.is_cuda(), "Input tensor must be a CUDA tensor");
    TORCH_CHECK(input.dim() == 3, "Input tensor must be 3D");
    TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous");

    const int N = input.size(0);
    const int C = input.size(1);
    const int L_in = input.size(2);
    
    const int L_out = floor(((float)L_in - kernel_size) / stride) + 1;
    
    auto output = torch::empty({N, C, L_out}, input.options());
    const int64_t n_elements_out = output.numel();
    
    if (n_elements_out == 0) {
        return output;
    }

    const int block_size = 256;
    const int num_blocks = (n_elements_out + block_size - 1) / block_size;

    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "lppool1d_kernel", ([&] {
        lppool1d_kernel<scalar_t><<<num_blocks, block_size>>>(
            output.data_ptr<scalar_t>(),
            input.data_ptr<scalar_t>(),
            static_cast<scalar_t>(norm_power),
            kernel_size, stride,
            N, C, L_in, L_out, n_elements_out
        );
    }));

    return output;
}
"""

class ModelNew(nn.Module):
    """
    使用自定义 CUDA 内核进行优化的 LPPool1d 模型。
    """
    def __init__(self, norm_power, kernel_size, stride):
        super(ModelNew, self).__init__()
        self.norm_power = norm_power
        self.kernel_size = kernel_size
        self.stride = stride
        
        self.op = load_inline(
            name='lppool1d_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['lppool1d_cuda_forward'],
            verbose=False
        )

    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
        return self.op.lppool1d_cuda_forward(
            input_tensor.contiguous(), # 确保输入是连续的
            self.norm_power,
            self.kernel_size,
            self.stride
        )