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

instancenorm_minkowski_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

// 线程块内归约求和
__inline__ __device__ float blockReduceSum(float val, int tid) {
    extern __shared__ float sdata[];
    sdata[tid] = val;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sdata[tid] += sdata[tid + stride];
        }
        __syncthreads();
    }
    return sdata[0];
}

// 融合内核 - InstanceNorm -> Minkowski
__global__ void instancenorm_minkowski_kernel(
    const float* __restrict__ x,
    float* __restrict__ y,
    int N, int C, int H, int W,
    float p, float inv_p, float eps
) {
    // blockIdx.x: sample index (N)
    // blockIdx.y: channel index (C)
    int sample_idx = blockIdx.x;
    int channel_idx = blockIdx.y;

    if (sample_idx >= N || channel_idx >= C) return;

    int tid = threadIdx.x;
    int spatial_size = H * W;
    
    // 计算当前样本和通道的起始偏移量
    int base_offset = sample_idx * C * spatial_size + channel_idx * spatial_size;

    // --- 融合步骤1：计算InstanceNorm所需的统计量 ---
    float local_sum = 0.0f;
    for (int i = tid; i < spatial_size; i += blockDim.x) {
        local_sum += x[base_offset + i];
    }
    float total_sum = blockReduceSum(local_sum, tid);
    __shared__ float mean_val;
    if (tid == 0) mean_val = total_sum / spatial_size;
    __syncthreads();

    float local_sum_sq = 0.0f;
    for (int i = tid; i < spatial_size; i += blockDim.x) {
        float val = x[base_offset + i] - mean_val;
        local_sum_sq += val * val;
    }
    float total_sum_sq = blockReduceSum(local_sum_sq, tid);
    __shared__ float inv_std_val;
    if (tid == 0) inv_std_val = rsqrtf(total_sum_sq / spatial_size + eps);
    __syncthreads();

    // --- 融合步骤2：计算归一化后特征的Minkowski范数 ---
    float local_minkowski_sum = 0.0f;
    for (int i = tid; i < spatial_size; i += blockDim.x) {
        float normed_val = (x[base_offset + i] - mean_val) * inv_std_val;
        local_minkowski_sum += powf(fabsf(normed_val), p);
    }
    float total_minkowski_sum = blockReduceSum(local_minkowski_sum, tid);
    
    // 第一个线程计算最终结果并写入输出
    if (tid == 0) {
        float minkowski_val = powf(total_minkowski_sum, inv_p);
        int output_offset = sample_idx * C + channel_idx;
        y[output_offset] = minkowski_val;
    }
}

torch::Tensor instancenorm_minkowski_cuda(torch::Tensor x, float p, float eps) {
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "X must be float32");
    TORCH_CHECK(x.dim() == 4, "X must be 4D");
    TORCH_CHECK(p > 0, "p must be positive");
    
    auto x_contig = x.contiguous();
    
    int N = x_contig.size(0);
    int C = x_contig.size(1);
    int H = x_contig.size(2);
    int W = x_contig.size(3);
    
    // 输出是 [N, C]
    auto y = torch::zeros({N, C}, x.options());
    
    float inv_p = 1.0f / p;
    const int block_size = 256;
    size_t shared_mem = block_size * sizeof(float);
    
    dim3 blocks(N, C);
    
    instancenorm_minkowski_kernel<<<blocks, block_size, shared_mem>>>(
        x_contig.data_ptr<float>(),
        y.data_ptr<float>(),
        N, C, H, W,
        p, inv_p, eps
    );
    
    return y;
}
"""

instancenorm_minkowski_cpp_source = """
torch::Tensor instancenorm_minkowski_cuda(torch::Tensor x, float p, float eps);
"""

instancenorm_minkowski = load_inline(
    name="instancenorm_minkowski",
    cpp_sources=instancenorm_minkowski_cpp_source,
    cuda_sources=instancenorm_minkowski_source,
    functions=["instancenorm_minkowski_cuda"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, num_features=64, p=2, eps=1e-5):
        super(ModelNew, self).__init__()
        self.p = p
        self.eps = eps
        self.instancenorm_minkowski = instancenorm_minkowski

    def forward(self, x):
        return self.instancenorm_minkowski.instancenorm_minkowski_cuda(x, self.p, self.eps)
