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

# -------------------------------------------------------------
# 常量定义 (同上)
# -------------------------------------------------------------
N, C, H, W = 8, 10, 16, 16
# 确保 weight 在 CUDA 上
WEIGHT = torch.rand(C, dtype=torch.float32).cuda()
IGNORE_INDEX = -100
REDUCTION = 'mean'
BLOCK_SIZE = 256


# -------------------------------------------------------------

class ModelNew(nn.Module):

    def __init__(self, weight=None, size_average=None, ignore_index=-100,
                 reduce=None, reduction='mean'):
        super().__init__()
        self.reduction_str = reduction
        self.ignore_index = ignore_index

        if weight is not None:
            self.register_buffer('weight', weight.contiguous())
        else:
            self.weight = None

        self.block_size = BLOCK_SIZE
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):

        cpp_header = f"""
        #include <torch/extension.h>

        // C++ 接口
        torch::Tensor nll_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            c10::optional<torch::Tensor> weight,
            int64_t ignore_index,
            std::string reduction
        );
        """

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

        #define BLOCK_SIZE {self.block_size}

        /* * 辅助函数：计算单个元素的 loss 和 weight
         * (在两个核函数之间共享)
         */
        __device__ inline void compute_nll_loss_item(
            const float* __restrict__ input_data,
            const int64_t* __restrict__ target_data, 
            const float* __restrict__ weight_data,
            int i, int N_spatial, int C, int64_t ignore_index,
            float& loss_val, // 输出
            float& weight_val // 输出
        ) {{
            const int64_t target_idx = target_data[i];

            weight_val = 1.0f;
            loss_val = 0.0f;

            if (target_idx == ignore_index) {{
                weight_val = 0.0f;
            }} else if (target_idx < 0 || target_idx >= C) {{
                weight_val = 0.0f;
            }} else {{
                if (weight_data != nullptr) {{
                    weight_val = weight_data[target_idx];
                }}

                const int n = i / N_spatial;
                const int s = i % N_spatial;
                // const int s = i - n * N_spatial; // 优化的 modulo

                const int64_t input_idx = 
                    (int64_t)n * C * N_spatial + 
                    (int64_t)target_idx * N_spatial + 
                    (int64_t)s;

                loss_val = -weight_val * input_data[input_idx];
            }}
        }}


        __global__ void nll_loss_kernel_no_reduce(
            const float* __restrict__ input_data,
            const int64_t* __restrict__ target_data,
            const float* __restrict__ weight_data,
            float* __restrict__ loss_out_data, 
            int N_spatial, int C, int N_total, int64_t ignore_index
        ) {{
            for (int i = blockIdx.x * blockDim.x + threadIdx.x; 
                 i < N_total; 
                 i += gridDim.x * blockDim.x) 
            {{
                float loss_val, weight_val;
                compute_nll_loss_item(
                    input_data, target_data, weight_data,
                    i, N_spatial, C, ignore_index,
                    loss_val, weight_val
                );
                // 写入未归约的损失
                loss_out_data[i] = loss_val;
            }}
        }}


        __global__ void nll_loss_kernel_reduce(
            const float* __restrict__ input_data,
            const int64_t* __restrict__ target_data,
            const float* __restrict__ weight_data,
            float* __restrict__ partial_loss_out,    // Block-level
            float* __restrict__ partial_weight_out,  // Block-level
            int N_spatial, int C, int N_total, int64_t ignore_index
        ) {{
            // 共享内存用于 Block 内部归约
            __shared__ float s_loss[BLOCK_SIZE];
            __shared__ float s_weight[BLOCK_SIZE];

            const int tid = threadIdx.x;
            float thread_loss_sum = 0.0f;
            float thread_weight_sum = 0.0f;

            // 1. Grid-Stride Loop: 计算和累加
            for (int i = blockIdx.x * blockDim.x + threadIdx.x; 
                 i < N_total; 
                 i += gridDim.x * blockDim.x) 
            {{
                float loss_val, weight_val;
                compute_nll_loss_item(
                    input_data, target_data, weight_data,
                    i, N_spatial, C, ignore_index,
                    loss_val, weight_val
                );
                thread_loss_sum += loss_val;
                thread_weight_sum += weight_val;
            }}

            s_loss[tid] = thread_loss_sum;
            s_weight[tid] = thread_weight_sum;
            __syncthreads();

            // 2. 共享内存归约 (Sum)
            for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {{
                if (tid < offset) {{
                    s_loss[tid] += s_loss[tid + offset];
                    s_weight[tid] += s_weight[tid + offset];
                }}
                __syncthreads();
            }}

            // 3. 线程 0 写入 Block 的总和
            if (tid == 0) {{
                partial_loss_out[blockIdx.x] = s_loss[0];
                partial_weight_out[blockIdx.x] = s_weight[0];
            }}
        }}

        // C++ 封装函数
        torch::Tensor nll_loss_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            c10::optional<torch::Tensor> weight_opt,
            int64_t ignore_index,
            std::string reduction
        ) {{
            // 检查
            TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
            TORCH_CHECK(target.is_cuda(), "target must be a CUDA tensor");
            TORCH_CHECK(input.dim() >= 2, "input must be >= 2D");
            TORCH_CHECK(input.dim() == target.dim() + 1, "input/target dim mismatch");
            TORCH_CHECK(target.scalar_type() == torch::kLong, "target must be torch.long");

            const int64_t N = input.size(0);
            const int64_t C = input.size(1);

            int64_t N_spatial = 1;
            if (input.dim() > 2) {{
                for (int d = 2; d < input.dim(); ++d) {{
                    TORCH_CHECK(input.size(d) == target.size(d-1), "spatial dim mismatch");
                    N_spatial *= input.size(d);
                }}
            }}
            const int64_t N_total = N * N_spatial;

            if (N_total == 0) {{
                return torch::tensor(0.0, input.options());
            }}

            auto input_flat = input.contiguous().view({{N, C, N_spatial}});
            auto target_flat = target.contiguous().view({{N_total}});

            const float* weight_ptr = nullptr;
            if (weight_opt.has_value()) {{
                auto& weight = weight_opt.value();
                TORCH_CHECK(weight.is_cuda() && weight.is_contiguous() && 
                            weight.dim() == 1 && weight.size(0) == C, "weight size mismatch");
                weight_ptr = weight.data_ptr<float>();
            }}

            // 4.核函数路由
            dim3 block_dim(BLOCK_SIZE);
            dim3 grid_dim((N_total + BLOCK_SIZE - 1) / BLOCK_SIZE);

            if (reduction == "none") {{
                auto output_loss = torch::empty_like(target_flat, input.options());

                nll_loss_kernel_no_reduce<<<grid_dim, block_dim>>>(
                    input_flat.data_ptr<float>(),
                    target_flat.data_ptr<long>(),
                    weight_ptr,
                    output_loss.data_ptr<float>(),
                    N_spatial, C, N_total, ignore_index
                );
                return output_loss.view(target.sizes());

            }} else {{ // "mean" or "sum"

                // 创建小的部分和张量
                auto partial_loss_out = torch::empty({{grid_dim.x}}, input.options());
                auto partial_weight_out = torch::empty({{grid_dim.x}}, input.options());

                nll_loss_kernel_reduce<<<grid_dim, block_dim>>>(
                    input_flat.data_ptr<float>(),
                    target_flat.data_ptr<long>(),
                    weight_ptr,
                    partial_loss_out.data_ptr<float>(),
                    partial_weight_out.data_ptr<float>(),
                    N_spatial, C, N_total, ignore_index
                );

                // 5. 在 C++ 中对*小的*部分和张量进行归约
                torch::Tensor total_loss = partial_loss_out.sum();

                if (reduction == "sum") {{
                    return total_loss;
                }} 

                // reduction == "mean"
                double total_weight = partial_weight_out.sum().item<double>();
                if (total_weight == 0.0) {{
                    return torch::tensor(0.0, input.options());
                }}
                return total_loss / total_weight;
            }}
        }}
        """

        nvcc_flags = ['-O3', '--use_fast_math']

        # JIT (Just-In-Time) 编译
        self.loss_op = load_inline(
            name="nll_loss_op_v2_reduce",  # 更改名称
            cpp_sources=cpp_header,
            cuda_sources=cuda_source,
            functions=["nll_loss_forward_cuda"],
            extra_cuda_cflags=nvcc_flags,
            verbose=False
        )

    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:

        # C++ wrapper 现在处理 .contiguous()
        return self.loss_op.nll_loss_forward_cuda(
            input,
            target,
            self.weight,
            self.ignore_index,
            self.reduction_str
        )