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

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

// 分块归约函数
__inline__ __device__ void blockReduceSum(float* sdata, float val, int tid) {
    sdata[tid] = val;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sdata[tid] += sdata[tid + stride];
        }
        __syncthreads();
    }
}

// InstanceNorm + ReLU融合kernel
__global__ void instancenorm_relu_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ y,
    int batch, int channels, int height, int width, float eps
) {
    int spatial_size = height * width;
    int instance_idx = blockIdx.x;
    int channel_idx = blockIdx.y;

    if (instance_idx >= batch || channel_idx >= channels) return;

    int tid = threadIdx.x;
    int instance_offset = instance_idx * channels * spatial_size + channel_idx * spatial_size;

    // 共享内存：sum和sum_sq
    extern __shared__ float shared_mem[];
    float* sum_smem = shared_mem;
    float* sum_sq_smem = shared_mem + blockDim.x;

    // 每个线程计算部分和
    float local_sum = 0.0f;
    float local_sum_sq = 0.0f;

    for (int i = tid; i < spatial_size; i += blockDim.x) {
        float val = x[instance_offset + i];
        local_sum += val;
        local_sum_sq += val * val;
    }

    // 归约求和
    blockReduceSum(sum_smem, local_sum, tid);
    blockReduceSum(sum_sq_smem, local_sum_sq, tid);

    // 计算统计量并广播
    __shared__ float mean_val;
    __shared__ float inv_std_val;
    __shared__ float weight_val;
    __shared__ float bias_val;

    if (tid == 0) {
        float mean = sum_smem[0] / spatial_size;
        float var = (sum_sq_smem[0] / spatial_size) - (mean * mean);
        var = fmaxf(var, 0.0f);

        mean_val = mean;
        inv_std_val = rsqrtf(var + eps);
        weight_val = weight[channel_idx];
        bias_val = bias[channel_idx];
    }
    __syncthreads();

    // 应用InstanceNorm + ReLU
    for (int i = tid; i < spatial_size; i += blockDim.x) {
        int idx = instance_offset + i;
        float val = x[idx];
        float normalized = (val - mean_val) * inv_std_val * weight_val + bias_val;
        
        // 融合ReLU：max(0, normalized)
        y[idx] = fmaxf(normalized, 0.0f);
    }
}

torch::Tensor instancenorm_relu_cuda_forward(
    torch::Tensor x, torch::Tensor weight, torch::Tensor bias, float eps
) {
    auto x_contig = x.contiguous();
    int batch = x_contig.size(0);
    int channels = x_contig.size(1);
    int height = x_contig.size(2);
    int width = x_contig.size(3);

    auto y = torch::empty_like(x_contig);

    dim3 blocks(batch, channels);
    int threads = 256;
    size_t shared_mem = 2 * threads * sizeof(float) + 4 * sizeof(float);

    instancenorm_relu_kernel<<<blocks, threads, shared_mem>>>(
        x_contig.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        y.data_ptr<float>(),
        batch, channels, height, width, eps
    );

    return y;
}
"""

instancenorm_relu_cpp_source = """
torch::Tensor instancenorm_relu_cuda_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, float eps);
"""

# 编译CUDA扩展
instancenorm_relu = load_inline(
    name="instancenorm_relu_fused",
    cpp_sources=instancenorm_relu_cpp_source,
    cuda_sources=instancenorm_relu_source,
    functions=["instancenorm_relu_cuda_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=True
)

class ModelNew(torch.nn.Module):
    """
    InstanceNorm + ReLU融合模型
    """
    def __init__(self, num_features=64, eps=1e-5, affine=True, track_running_stats=False):
        super(ModelNew, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.affine = affine
        self.track_running_stats = False  # 强制为False

        if affine:
            self.weight = torch.nn.Parameter(torch.ones(num_features))
            self.bias = torch.nn.Parameter(torch.zeros(num_features))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        融合实现：InstanceNorm + ReLU
        
        Args:
            x (torch.Tensor): 输入张量 shape [B, C, H, W]
        
        Returns:
            torch.Tensor: ReLU(InstanceNorm(x))
        """
        if self.affine:
            return instancenorm_relu.instancenorm_relu_cuda_forward(
                x, self.weight, self.bias, self.eps
            )
        else:
            weight = torch.ones(self.num_features, device=x.device)
            bias = torch.zeros(self.num_features, device=x.device)
            return instancenorm_relu.instancenorm_relu_cuda_forward(
                x, weight, bias, self.eps
            )
