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

instancenorm_relu_dropout_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 + Dropout三重融合kernel
__global__ void instancenorm_relu_dropout_kernel(
    const float* __restrict__ x,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    const float* __restrict__ dropout_mask,  // 预生成的dropout mask
    float* __restrict__ y,
    int batch, int channels, int height, int width, 
    float eps, float dropout_scale, bool training
) {
    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 + Dropout
    for (int i = tid; i < spatial_size; i += blockDim.x) {
        int idx = instance_offset + i;
        int mask_idx = instance_idx * channels * spatial_size + channel_idx * spatial_size + i;
        
        float val = x[idx];
        float normalized = (val - mean_val) * inv_std_val * weight_val + bias_val;
        
        // 融合ReLU：max(0, normalized)
        float activated = fmaxf(normalized, 0.0f);
        
        if (training) {
            // 使用预生成的mask：0表示丢弃，dropout_scale表示保留并缩放
            y[idx] = activated * dropout_mask[mask_idx];
        } else {
            // 推理模式：只应用InstanceNorm + ReLU
            y[idx] = activated;
        }
    }
}

torch::Tensor instancenorm_relu_dropout_cuda_forward(
    torch::Tensor x, torch::Tensor weight, torch::Tensor bias, 
    torch::Tensor dropout_mask, float eps, float dropout_scale, bool training
) {
    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_dropout_kernel<<<blocks, threads, shared_mem>>>(
        x_contig.data_ptr<float>(),
        weight.data_ptr<float>(),
        bias.data_ptr<float>(),
        dropout_mask.data_ptr<float>(),
        y.data_ptr<float>(),
        batch, channels, height, width, eps, dropout_scale, training
    );

    return y;
}
"""

instancenorm_relu_dropout_cpp_source = """
torch::Tensor instancenorm_relu_dropout_cuda_forward(
    torch::Tensor x, torch::Tensor weight, torch::Tensor bias, 
    torch::Tensor dropout_mask, float eps, float dropout_scale, bool training
);
"""

# 编译CUDA扩展
instancenorm_relu_dropout = load_inline(
    name="instancenorm_relu_dropout_fused",
    cpp_sources=instancenorm_relu_dropout_cpp_source,
    cuda_sources=instancenorm_relu_dropout_source,
    functions=["instancenorm_relu_dropout_cuda_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=True
)

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

        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 + Dropout
        
        Args:
            x (torch.Tensor): 输入张量 shape [B, C, H, W]
        
        Returns:
            torch.Tensor: Dropout(ReLU(InstanceNorm(x)))
        """
        if self.training:
            # 使用PyTorch的Dropout生成mask，确保完全一致
            with torch.no_grad():
                # 创建一个与x相同的张量用于生成mask
                dummy_input = torch.ones_like(x)
                dropout_mask = torch.nn.functional.dropout(dummy_input, p=self.dropout_p, training=True, inplace=False)
                # 将mask转换为0或dropout_scale的形式
                dropout_mask = dropout_mask / dummy_input
        else:
            # 推理时mask全为1
            dropout_mask = torch.ones_like(x)
        
        if self.affine:
            return instancenorm_relu_dropout.instancenorm_relu_dropout_cuda_forward(
                x, self.weight, self.bias, dropout_mask, self.eps, self.dropout_scale, self.training
            )
        else:
            weight = torch.ones(self.num_features, device=x.device)
            bias = torch.zeros(self.num_features, device=x.device)
            return instancenorm_relu_dropout.instancenorm_relu_dropout_cuda_forward(
                x, weight, bias, dropout_mask, self.eps, self.dropout_scale, self.training
            )
