import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
# 从 torch 文件导入常量
from infonceloss_torch import BATCH_SIZE, FEATURE_DIM, TEMPERATURE, N_NEGATIVES


class ModelNew(nn.Module):

    def __init__(self):
        super().__init__()
        self.temperature = TEMPERATURE
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>

        // C++ 接口
        torch::Tensor infonce_forward_cuda(
            torch::Tensor query,              // (B, D)
            torch::Tensor positive,           // (B, D)
            torch::Tensor negative_sims,      // (B, N) - 预先计算的
            float temperature
        );
        """

        cuda_source = """
        #include <torch/extension.h> 
        #include <cuda_runtime.h>
        #include <cmath>
        #include <float.h> // For FLT_MAX

        // 使用 256 个线程的块大小
        #define BLOCK_SIZE 256

        /*
         * InfoNCE 融合核函数
         * 我们启动 B 个块 (gridDim.x = B)，每个块负责一行 (一个 query) 的 loss 计算。
         * 每个块 (blockIdx.x) 计算:
         * 1. query[i] 和 positive[i] 之间的点积 (pos_logit)
         * 2. 对 [pos_logit, neg_logits[i,:]] 执行稳定的 LogSumExp
         * 3. 计算 loss_i = -pos_logit + logsumexp
         *
         * @param loss_per_row_out - (B,) 形状的张量，用于存储 loss_i
         */
        __global__ void infonce_fused_kernel(
            const float* __restrict__ query_data,        // (B, D)
            const float* __restrict__ positive_data,     // (B, D)
            const float* __restrict__ negative_sims_data, // (B, N)
            float* __restrict__ loss_per_row_out, // (B,)
            int B,
            int D,
            int N,
            float temperature
        ) {
            // 每个块计算一行
            int i = blockIdx.x; // 当前 query 的索引 (0 到 B-1)
            if (i >= B) return; 

            // --- 共享内存 ---
            // s_dot 用于计算 pos_logit
            __shared__ float s_dot[BLOCK_SIZE];
            // s_max 和 s_sum 用于稳定的 LogSumExp
            __shared__ float s_max[BLOCK_SIZE];
            __shared__ float s_sum[BLOCK_SIZE];

            // --- 1. 计算 Positive Logit ---
            // 融合了 F.cosine_similarity(query[i], positive[i]) / temp
            float thread_dot_sum = 0.0f;

            // 使用 Grid-Stride 循环计算点积
            for (int k = threadIdx.x; k < D; k += BLOCK_SIZE) {
                thread_dot_sum += query_data[i * D + k] * positive_data[i * D + k];
            }
            s_dot[threadIdx.x] = thread_dot_sum;

            // 块内归约 (Sum)
            __syncthreads();
            for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
                if (threadIdx.x < offset) {
                    s_dot[threadIdx.x] += s_dot[threadIdx.x + offset];
                }
                __syncthreads();
            }

            // 线程 0 现在拥有 pos_logit
            // 我们将其存储在 s_dot[0] 中以供后续步骤使用
            if (threadIdx.x == 0) {
                s_dot[0] = s_dot[0] / temperature;
            }
            __syncthreads(); // 确保所有线程都能读到 s_dot[0]

            const float pos_logit = s_dot[0]; // 所有线程的常量

            // --- 2. 稳定的 LogSumExp (Pass 1: Find Max) ---
            float thread_max = -FLT_MAX;

            // 线程 0 包含 pos_logit
            if (threadIdx.x == 0) {
                thread_max = pos_logit;
            }

            // 遍历 N 个 negative logits
            for (int j = threadIdx.x; j < N; j += BLOCK_SIZE) {
                float neg_logit = negative_sims_data[i * N + j] / temperature;
                thread_max = fmaxf(thread_max, neg_logit);
            }
            s_max[threadIdx.x] = thread_max;

            // 块内归约 (Max)
            __syncthreads();
            for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
                if (threadIdx.x < offset) {
                    s_max[threadIdx.x] = fmaxf(s_max[threadIdx.x], s_max[threadIdx.x + offset]);
                }
                __syncthreads();
            }

            // 线程 0 拥有 global_max
            if (threadIdx.x == 0) {
                s_max[0] = s_max[0];
            }
            __syncthreads(); // 确保所有线程都能读到 s_max[0]

            const float global_max = s_max[0];

            // --- 3. 稳定的 LogSumExp (Pass 2: Sum Exp Diff) ---
            float thread_sum_exp = 0.0f;

            // 线程 0 添加 positive_logit 的贡献
            if (threadIdx.x == 0) {
                thread_sum_exp = expf(pos_logit - global_max);
            }

            // 遍历 N 个 negative logits
            for (int j = threadIdx.x; j < N; j += BLOCK_SIZE) {
                float neg_logit = negative_sims_data[i * N + j] / temperature;
                thread_sum_exp += expf(neg_logit - global_max);
            }
            s_sum[threadIdx.x] = thread_sum_exp;

            // 块内归约 (Sum)
            __syncthreads();
            for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
                if (threadIdx.x < offset) {
                    s_sum[threadIdx.x] += s_sum[threadIdx.x + offset];
                }
                __syncthreads();
            }

            // --- 4. 计算最终的 loss[i] ---
            if (threadIdx.x == 0) {
                float log_sum_exp = global_max + logf(s_sum[0]);
                // loss_i = -logit[0] + logsumexp
                float loss_i = -pos_logit + log_sum_exp;
                loss_per_row_out[i] = loss_i;
            }
        }

        // C++ 封装函数
        torch::Tensor infonce_forward_cuda(
            torch::Tensor query,
            torch::Tensor positive,
            torch::Tensor negative_sims, // 注意：这是未缩放的
            float temperature
        ) {
            // 检查
            TORCH_CHECK(query.is_cuda(), "query must be a CUDA tensor");
            TORCH_CHECK(positive.is_cuda(), "positive must be a CUDA tensor");
            TORCH_CHECK(negative_sims.is_cuda(), "negative_sims must be a CUDA tensor");

            query = query.contiguous();
            positive = positive.contiguous();
            negative_sims = negative_sims.contiguous();

            const int B = query.size(0);
            const int D = query.size(1);
            const int N = negative_sims.size(1);

            TORCH_CHECK(positive.size(0) == B && positive.size(1) == D, "positive tensor has wrong size");
            TORCH_CHECK(negative_sims.size(0) == B, "negative_sims tensor has wrong size");

            // 分配一个张量来保存每个块 (每行) 的 loss
            auto loss_per_row = torch::empty({B}, query.options());

            const int block_size = BLOCK_SIZE;
            const int grid_size = B; // B 个块，每个块处理一行

            // 启动 CUDA 核函数
            infonce_fused_kernel<<<grid_size, block_size>>>(
                query.data_ptr<float>(),
                positive.data_ptr<float>(),
                negative_sims.data_ptr<float>(),
                loss_per_row.data_ptr<float>(),
                B, D, N,
                temperature
            );


            // 核函数返回后，loss_per_row 包含 B 个 loss 值
            // 我们需要对它们取平均
            return loss_per_row.mean();
        }
        """

        # JIT (Just-In-Time) 编译
        self.infonce_op = load_inline(
            name="infonce_op_v1_stable",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["infonce_forward_cuda"],
            extra_cuda_cflags=["-O3"],
            verbose=False
        )

    def forward(self, query: torch.Tensor, positive: torch.Tensor, negatives: torch.Tensor) -> torch.Tensor:
        # 1. (Python) 执行优化的 matmul (cuBLAS)
        # (B, D) @ (D, N) -> (B, N)
        # 这是未缩放的 (没有 / temp)
        negative_sims_unscaled = torch.matmul(query, negatives.t())

        # 2. (CUDA) 调用融合核函数
        # 核函数将处理:
        # - query, positive 的 cosine similarity
        # - 对所有 sim 应用 / temp
        # - 稳定的 LogSumExp 和 CrossEntropy
        # - 最终的 Mean 归约
        return self.infonce_op.infonce_forward_cuda(
            query,
            positive,
            negative_sims_unscaled,
            self.temperature
        )