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

BATCH_SIZE = 4096
N_CLASSES = 1024


class ModelNew(nn.Module):

    def __init__(self):
        super().__init__()
        self._compile_cuda_kernel()

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

        // C++ 接口
        torch::Tensor cross_entropy_forward_cuda(
            torch::Tensor logits,
            torch::Tensor target
        );
        """

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

        #define BLOCK_SIZE 256

        __global__ void cross_entropy_fused_kernel(
            const float* __restrict__ logits_data,   // (N, C)
            const int64_t* __restrict__ target_data, // (N,)
            float* __restrict__ loss_per_row_out,    // (N,)
            int N,
            int C
        ) {
            // 当前处理的 Batch 索引
            int row_idx = blockIdx.x;
            if (row_idx >= N) return;

            // 当前行的指针
            const float* row_logits = logits_data + row_idx * C;
            int tid = threadIdx.x;

            // 共享内存：用于 Max 和 Sum 的归约
            __shared__ float s_data[BLOCK_SIZE];

            float thread_max = -FLT_MAX;

            // Grid-Stride Loop 遍历类别 C
            for (int c = tid; c < C; c += BLOCK_SIZE) {
                float val = row_logits[c];
                if (val > thread_max) {
                    thread_max = val;
                }
            }
            s_data[tid] = thread_max;
            __syncthreads();

            // 块内归约 (Max)
            for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {
                if (tid < offset) {
                    if (s_data[tid + offset] > s_data[tid]) {
                        s_data[tid] = s_data[tid + offset];
                    }
                }
                __syncthreads();
            }


            float row_max_val = s_data[0];
            __syncthreads(); 

            float thread_sum_exp = 0.0f;

            for (int c = tid; c < C; c += BLOCK_SIZE) {
                float val = row_logits[c];
                thread_sum_exp += expf(val - row_max_val);
            }
            s_data[tid] = thread_sum_exp;
            __syncthreads();

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

            if (tid == 0) {
                float row_sum_exp = s_data[0];
                float log_sum_exp = logf(row_sum_exp) + row_max_val;

                int64_t target_class = target_data[row_idx];
                float target_logit = row_logits[target_class];

                // Cross Entropy Formula
                loss_per_row_out[row_idx] = -target_logit + log_sum_exp;
            }
        }

        // C++ 封装函数
        torch::Tensor cross_entropy_forward_cuda(
            torch::Tensor logits,
            torch::Tensor target
        ) {
            TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
            TORCH_CHECK(target.is_cuda(), "target must be a CUDA tensor");
            TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
            TORCH_CHECK(target.dim() == 1, "target must be 1D");

            // 确保连续
            logits = logits.contiguous();
            target = target.contiguous();

            int N = logits.size(0); // Batch Size
            int C = logits.size(1); // Num Classes

            TORCH_CHECK(target.size(0) == N, "Target size mismatch");

            auto losses = torch::empty({N}, logits.options());

            // 启动配置: 
            // Grid: N (每个 Batch 一个 Block)
            // Block: 256
            dim3 grid_dim(N);
            dim3 block_dim(BLOCK_SIZE);

            cross_entropy_fused_kernel<<<grid_dim, block_dim>>>(
                logits.data_ptr<float>(),
                target.data_ptr<int64_t>(),
                losses.data_ptr<float>(),
                N, C
            );

            // 返回 Mean Reduction
            return losses.mean();
        }
        """

        self.ce_op = load_inline(
            name="cross_entropy_op_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["cross_entropy_forward_cuda"],
            extra_cuda_cflags=["-O3"],
            verbose=False
        )

    def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return self.ce_op.cross_entropy_forward_cuda(logits, target)