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

# 定义维度常量
N, C, H, W = 32, 64, 56, 56
EPS = 1e-6
INSTANCE_SIZE = C * H * W  # 200704

# 确保 C*H*W 是 4 的倍数以便向量化
assert (INSTANCE_SIZE) % 4 == 0, "Instance size (C*H*W) must be a multiple of 4"


class ModelNew(nn.Module):
    """
    Tanimoto 系数的 CUDA 优化实现。

    优化点:
    1. 融合: 将 3 次归约 (dot, x_sq, y_sq) 和 1 次除法融合到一个 CUDA kernel。
    2. 向量化: 使用 float4 进行 128 位内存访问。
    3. 共享内存: 用于高效的块内归约。
    4. 精度匹配: 使用 'volatile' 禁用 FMA，以 100% 匹配 PyTorch 的分步舍入。
    """

    def __init__(self):
        super().__init__()
        self.eps = EPS
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        # 1. C++ 接口定义
        cpp_source = """
        #include <torch/extension.h>

        torch::Tensor tanimoto_forward_cuda(
            torch::Tensor x,
            torch::Tensor y,
            float eps,
            int N,
            int D);
        """

        # 2. CUDA 内核实现
        cuda_source = f"""
        #include <cuda_runtime.h>
        #include <device_launch_parameters.h>
        #include <cmath>

        // ============================================================
        // Tanimoto 融合内核
        // ============================================================
        __global__ void tanimoto_fused_kernel(
            const float* __restrict__ x,
            const float* __restrict__ y,
            float* __restrict__ output,
            float eps,
            int N,          // 批次大小
            int D           // 实例大小 (C*H*W)
        ) {{
            // 每个块处理一个实例 (n_idx)
            const int n_idx = blockIdx.x;
            if (n_idx >= N) return;

            const int t_idx = threadIdx.x;
            const int block_size = blockDim.x; // e.g., 256

            // 指向此实例数据的开头
            const int instance_offset = n_idx * D;
            const float4* x4_ptr = reinterpret_cast<const float4*>(x + instance_offset);
            const float4* y4_ptr = reinterpret_cast<const float4*>(y + instance_offset);

            const int D_vec = D / 4; // 向量化维度

            // 线程本地累加器
            float thread_dot = 0.0f;
            float thread_x_sq = 0.0f;
            float thread_y_sq = 0.0f;

            // --- 1. 线程本地向量化累加 ---
            for (int i = t_idx; i < D_vec; i += block_size) {{
                float4 x_val = x4_ptr[i];
                float4 y_val = y4_ptr[i];

                // --- 精度修复: 禁用 FMA ---
                // (x * y)
                volatile float m_dot_x = x_val.x * y_val.x;
                volatile float m_dot_y = x_val.y * y_val.y;
                volatile float m_dot_z = x_val.z * y_val.z;
                volatile float m_dot_w = x_val.w * y_val.w;

                // (x * x)
                volatile float m_x_sq_x = x_val.x * x_val.x;
                volatile float m_x_sq_y = x_val.y * x_val.y;
                volatile float m_x_sq_z = x_val.z * x_val.z;
                volatile float m_x_sq_w = x_val.w * x_val.w;

                // (y * y)
                volatile float m_y_sq_x = y_val.x * y_val.x;
                volatile float m_y_sq_y = y_val.y * y_val.y;
                volatile float m_y_sq_z = y_val.z * y_val.z;
                volatile float m_y_sq_w = y_val.w * y_val.w;

                thread_dot += m_dot_x + m_dot_y + m_dot_z + m_dot_w;
                thread_x_sq += m_x_sq_x + m_x_sq_y + m_x_sq_z + m_x_sq_w;
                thread_y_sq += m_y_sq_x + m_y_sq_y + m_y_sq_z + m_y_sq_w;
            }}

            // --- 2. 块内共享内存归约 ---
            // (BLOCK_SIZE 必须是 2 的幂, e.g., 256)
            extern __shared__ float s_data[];
            float* s_dot = s_data;
            float* s_x_sq = &s_data[block_size];
            float* s_y_sq = &s_data[block_size * 2];

            s_dot[t_idx] = thread_dot;
            s_x_sq[t_idx] = thread_x_sq;
            s_y_sq[t_idx] = thread_y_sq;

            __syncthreads();

            // (使用标准树形归约)
            for (int s = block_size / 2; s > 0; s >>= 1) {{
                if (t_idx < s) {{
                    s_dot[t_idx] += s_dot[t_idx + s];
                    s_x_sq[t_idx] += s_x_sq[t_idx + s];
                    s_y_sq[t_idx] += s_y_sq[t_idx + s];
                }}
                __syncthreads();
            }}

            // --- 3. 线程 0 计算最终结果 ---
            if (t_idx == 0) {{
                float total_dot = s_dot[0];
                float total_x_sq = s_x_sq[0];
                float total_y_sq = s_y_sq[0];

                float denominator = total_x_sq + total_y_sq - total_dot + eps;

                output[n_idx] = (total_dot + eps) / denominator;
            }}
        }}

        // ============================================================
        // C++ Wrapper
        // ============================================================
        torch::Tensor tanimoto_forward_cuda(
            torch::Tensor x,
            torch::Tensor y,
            float eps,
            int N,
            int D          // 实例大小 (C*H*W)
        ) {{
            // 确保输入连续
            x = x.contiguous();
            y = y.contiguous();

            // 输出张量形状为 [N]
            auto output = torch::empty({{N}}, x.options());

            const int BLOCK_SIZE = 256;
            dim3 threads(BLOCK_SIZE);
            dim3 blocks(N);

            // 动态分配共享内存 (每个累加器 256 * 4 字节)
            // (dot_sum, x_sq_sum, y_sq_sum)
            int shared_mem_size = 3 * BLOCK_SIZE * sizeof(float);

            tanimoto_fused_kernel<<<blocks, threads, shared_mem_size>>>(
                x.data_ptr<float>(),
                y.data_ptr<float>(),
                output.data_ptr<float>(),
                eps,
                N,
                D
            );

            return output;
        }}
        """

        # 3. 编译 CUDA 模块
        self.tanimoto_op = load_inline(
            name="tanimoto_cuda_v1_fma_fixed",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["tanimoto_forward_cuda"],
            # (使用高精度标志, 禁用 fast_math 以确保 volatile 生效)
            extra_cuda_cflags=[
                "-O3",
                "-ftz=false",
                "-prec-div=true",
                "-prec-sqrt=true",
                "-std=c++17"
            ],
            verbose=False
        )

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        if x.dtype != torch.float32 or not x.is_cuda:
            x = x.to("cuda", dtype=torch.float32)
        if y.dtype != torch.float32 or not y.is_cuda:
            y = y.to("cuda", dtype=torch.float32)

        N, C, H, W = x.size()
        D = C * H * W  # 实例大小

        return self.tanimoto_op.tanimoto_forward_cuda(
            x.contiguous(),
            y.contiguous(),
            self.eps,
            N,
            D
        )

