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

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

const int WARP_SIZE = 32;

// 优化的warp级归约
__inline__ __device__ float warpReduceSum(float val) {
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__inline__ __device__ float warpReduceSumSq(float val) {
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

// 优化的InstanceNorm内核 - 使用warp级和block级混合归约
__global__ void instancenorm_optimized_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 lane_id = tid % WARP_SIZE;
    int warp_id = tid / WARP_SIZE;
    int instance_offset = instance_idx * channels * spatial_size + channel_idx * spatial_size;

    extern __shared__ float sdata[];
    float* warp_sums = sdata;
    float* warp_sum_sqs = sdata + (blockDim.x / WARP_SIZE) * 2;

    // 第一阶段：每个warp内部归约
    float sum = 0.0f;
    float sum_sq = 0.0f;

    // 使用循环展开和向量化友好的访问模式
    for (int i = tid; i < spatial_size; i += blockDim.x) {
        float v = x[instance_offset + i];
        sum += v;
        sum_sq += v * v;
    }

    // Warp级归约
    sum = warpReduceSum(sum);
    sum_sq = warpReduceSumSq(sum_sq);

    // 每个warp的第一个线程保存结果到shared memory
    if (lane_id == 0) {
        warp_sums[warp_id] = sum;
        warp_sum_sqs[warp_id] = sum_sq;
    }
    __syncthreads();

    // 第二阶段：block级归约（只在warp 0中进行）
    if (warp_id == 0) {
        sum = (lane_id < (blockDim.x / WARP_SIZE)) ? warp_sums[lane_id] : 0.0f;
        sum_sq = (lane_id < (blockDim.x / WARP_SIZE)) ? warp_sum_sqs[lane_id] : 0.0f;

        // 再次warp归约
        sum = warpReduceSum(sum);
        sum_sq = warpReduceSumSq(sum_sq);

        // 计算最终统计量
        if (lane_id == 0) {
            float mean = sum / spatial_size;
            float var = (sum_sq / spatial_size) - (mean * mean);
            var = fmaxf(var, 0.0f);

            // 保存到shared memory供所有线程使用
            warp_sums[0] = mean;
            warp_sum_sqs[0] = rsqrtf(var + eps);
            warp_sums[1] = weight[channel_idx];
            warp_sum_sqs[1] = bias[channel_idx];
        }
    }
    __syncthreads();

    // 所有线程读取统计量
    float mean = warp_sums[0];
    float inv_std = warp_sum_sqs[0];
    float w = warp_sums[1];
    float b = warp_sum_sqs[1];

    // 应用InstanceNorm - 使用更优化的内存访问模式
    for (int i = tid; i < spatial_size; i += blockDim.x) {
        float v = x[instance_offset + i];
        float norm_val = (v - mean) * inv_std;
        y[instance_offset + i] = norm_val * w + b;
    }
}

// 针对小尺寸的优化内核
__global__ void instancenorm_small_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;

    extern __shared__ float sdata[];
    float* sum_shared = sdata;
    float* sum_sq_shared = sdata + blockDim.x;

    // 针对小尺寸的简化归约
    float sum = 0.0f;
    float sum_sq = 0.0f;

    for (int i = tid; i < spatial_size; i += blockDim.x) {
        float v = x[instance_offset + i];
        sum += v;
        sum_sq += v * v;
    }

    sum_shared[tid] = sum;
    sum_sq_shared[tid] = sum_sq;
    __syncthreads();

    // 树状归约
    for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
        if (tid < offset) {
            sum_shared[tid] += sum_shared[tid + offset];
            sum_sq_shared[tid] += sum_sq_shared[tid + offset];
        }
        __syncthreads();
    }

    __shared__ float s_mean;
    __shared__ float s_inv_std;
    __shared__ float s_weight;
    __shared__ float s_bias;

    if (tid == 0) {
        float mean = sum_shared[0] / spatial_size;
        float var = (sum_sq_shared[0] / spatial_size) - (mean * mean);
        var = fmaxf(var, 0.0f);
        s_mean = mean;
        s_inv_std = rsqrtf(var + eps);
        s_weight = weight[channel_idx];
        s_bias = bias[channel_idx];
    }
    __syncthreads();

    float mean = s_mean;
    float inv_std = s_inv_std;
    float w = s_weight;
    float b = s_bias;

    // 应用归一化
    for (int i = tid; i < spatial_size; i += blockDim.x) {
        float v = x[instance_offset + i];
        float norm_val = (v - mean) * inv_std;
        y[instance_offset + i] = norm_val * w + b;
    }
}

torch::Tensor instancenorm_cuda_forward(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    float eps
) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(weight.is_cuda(), "weight must be a CUDA tensor");
    TORCH_CHECK(bias.is_cuda(), "bias must be a CUDA tensor");
    TORCH_CHECK(x.dim() == 4, "input must be 4D [batch, channels, height, width]");
    TORCH_CHECK(weight.dim() == 1, "weight must be 1D");
    TORCH_CHECK(bias.dim() == 1, "bias must be 1D");
    TORCH_CHECK(x.size(1) == weight.size(0), "channel size mismatch");
    TORCH_CHECK(weight.size(0) == bias.size(0), "weight and bias size mismatch");

    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);
    int spatial_size = height * width;

    auto y = torch::empty_like(x_contig);

    // 更智能的线程配置
    dim3 blocks(batch, channels);
    size_t shared_mem;

    if (spatial_size >= 1024) {
        // 大尺寸使用优化内核，256线程，4个warp
        int threads = 256;
        shared_mem = (threads / WARP_SIZE) * 2 * sizeof(float) + 4 * sizeof(float);
        instancenorm_optimized_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
        );
    } else {
        // 小尺寸使用简化内核
        int threads;
        if (spatial_size <= 64) threads = 64;
        else if (spatial_size <= 128) threads = 128;
        else threads = 256;

        threads = min(threads, spatial_size);
        if (threads < 32) threads = 32;

        shared_mem = 2 * threads * sizeof(float);
        instancenorm_small_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
        );
    }

    // 移除同步，让CUDA流自动管理
    // cudaDeviceSynchronize(); 
    return y;
}
"""

instancenorm_cpp_source = """
torch::Tensor instancenorm_cuda_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias, float eps);
"""

instancenorm_cuda = load_inline(
    name="instancenorm_cuda",
    cpp_sources=instancenorm_cpp_source,
    cuda_sources=instancenorm_source,
    functions=["instancenorm_cuda_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=True
)


# CUDA优化版本 - 完全与PyTorch一致
class ModelNew(nn.Module):
    """
    Simplified CUDA version that forces track_running_stats=False for exact equivalence.
    """

    def __init__(self, num_features=64, eps=1e-5, affine=True, track_running_stats=False):
        super(ModelNew, self).__init__()

        # 强制track_running_stats=False以确保与CUDA实现完全等价
        if track_running_stats:
            print("警告：CUDA优化版本不支持track_running_stats=True，已强制设置为False")

        self.num_features = num_features
        self.eps = eps
        self.affine = affine
        self.track_running_stats = False  # 强制为False

        if affine:
            self.weight = nn.Parameter(torch.ones(num_features))
            self.bias = 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:
        """
        只支持track_running_stats=False的情况
        """
        if self.affine:
            return instancenorm_cuda.instancenorm_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_cuda.instancenorm_cuda_forward(x, weight, bias, self.eps)