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

# C++ 源代码 wrapper
cpp_source = """
#include <torch/extension.h>
#include <string>

torch::Tensor poly1_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& labels,
    float epsilon,
    std::string reduction);
"""

# CUDA 源代码
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>
#include <cfloat>

// 针对 C500 优化
#define WARP_SIZE 64
#define BLOCK_SIZE 256
#define ILP 4

struct __align__(16) Float4 {
    float x, y, z, w;
};

// Online Softmax 状态结构体
struct MD {
    float m; // max
    float d; // denominator (sum of exps)
};

// 合并状态
__device__ __forceinline__ MD reduce_md_op(MD a, MD b) {
    bool a_bigger = (a.m > b.m);
    float bigger_m = a_bigger ? a.m : b.m;
    float smaller_m = a_bigger ? b.m : a.m;
    MD res;
    res.m = bigger_m;
    if (res.m == -FLT_MAX) {
        res.d = 0.0f;
    } else {
        float diff = smaller_m - bigger_m;
        res.d = (a_bigger ? a.d : b.d) + (a_bigger ? b.d : a.d) * expf(diff);
    }
    return res;
}

// Warp Reduce
__device__ __forceinline__ MD warp_reduce_md(MD val) {
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        float other_m = __shfl_down_sync(0xffffffffffffffff, val.m, offset);
        float other_d = __shfl_down_sync(0xffffffffffffffff, val.d, offset);
        val = reduce_md_op(val, {other_m, other_d});
    }
    return val;
}

// Block Reduce
__device__ __forceinline__ MD block_reduce_md(MD val) {
    static __shared__ float shared_m[32]; 
    static __shared__ float shared_d[32];
    
    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;

    val = warp_reduce_md(val);

    if (lane == 0) {
        shared_m[wid] = val.m;
        shared_d[wid] = val.d;
    }
    __syncthreads();

    MD block_val = {-FLT_MAX, 0.0f};
    int num_warps = blockDim.x / WARP_SIZE;
    if (threadIdx.x < num_warps) {
        block_val.m = shared_m[threadIdx.x];
        block_val.d = shared_d[threadIdx.x];
    }
    
    if (wid == 0) {
        block_val = warp_reduce_md(block_val);
    }
    return block_val;
}

__global__ void poly1_loss_kernel(
    float* __restrict__ output,
    const float* __restrict__ logits,
    const int64_t* __restrict__ labels,
    int cols,
    float epsilon)
{
    int row_idx = blockIdx.x;
    int tid = threadIdx.x;

    const float* row_logits = logits + row_idx * cols;
    
    // 读取 Target Label
    // 只要有一个线程读取即可，这里让 tid 0 读取并广播，或者所有线程读取（如果都用得到）
    // 由于我们需要在最后一步计算 target_logit，所有线程可能并不需要知道 target_label，
    // 只有负责处理包含 target_label 那一列数据的线程需要吗？
    // 不，因为我们需要先算全行的 LogSumExp，这需要遍历所有数据。
    // 算出 LogSumExp 后，我们需要取出 row_logits[target_label]。
    // 既然所有线程已经遍历了数据，我们可以不用显式保存 target_logit，
    // 而是直接让线程 0 再读一次 target_logit (它就在 Global Memory)。
    // 或者在遍历时顺便把它抓出来存到 Shared Memory。
    // 考虑到 ILP 和 分块，直接让线程 0 最后读一次 target_logit 是最简单的，
    // 因为 target_logit 只是 1 个 float，额外的 1 次 global read 开销极小。

    // --- 1. 寄存器缓存 & Online Softmax ---
    Float4 reg_cache[ILP]; 
    MD state = {-FLT_MAX, 0.0f}; 

    int idx = tid * 4;
    #pragma unroll
    for (int k = 0; k < ILP; ++k) {
        if (idx < cols - 3) {
             reg_cache[k] = *(reinterpret_cast<const Float4*>(&row_logits[idx]));
             state = reduce_md_op(state, {reg_cache[k].x, 1.0f});
             state = reduce_md_op(state, {reg_cache[k].y, 1.0f});
             state = reduce_md_op(state, {reg_cache[k].z, 1.0f});
             state = reduce_md_op(state, {reg_cache[k].w, 1.0f});
        } else {
            // 边界处理：只读取有效数据
            if (idx < cols) {
                float v = row_logits[idx];
                state = reduce_md_op(state, {v, 1.0f});
            }
            if (idx + 1 < cols) {
                float v = row_logits[idx+1];
                state = reduce_md_op(state, {v, 1.0f});
            }
            if (idx + 2 < cols) {
                float v = row_logits[idx+2];
                state = reduce_md_op(state, {v, 1.0f});
            }
            if (idx + 3 < cols) {
                float v = row_logits[idx+3];
                state = reduce_md_op(state, {v, 1.0f});
            }
        }
        idx += BLOCK_SIZE * 4;
    }

    // --- 2. Block Reduction ---
    state = block_reduce_md(state);

    // --- 3. Final Calculation (Thread 0 only) ---
    if (tid == 0) {
        float global_max = state.m;
        float global_sum = state.d;
        float log_sum = global_max + logf(global_sum); // LogSumExp

        int64_t target_idx = labels[row_idx];
        
        // 直接读取 Target Logit
        float target_logit = row_logits[target_idx];
        
        // log_pt = x_t - LogSumExp
        float log_pt = target_logit - log_sum;
        float pt = expf(log_pt);
        
        // Loss = CE + eps * (1 - Pt) = -log_pt + eps * (1 - pt)
        float loss = -log_pt + epsilon * (1.0f - pt);
        
        output[row_idx] = loss;
    }
}

torch::Tensor poly1_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& labels, 
    float epsilon, 
    std::string reduction) 
{
    TORCH_CHECK(logits.is_cuda() && labels.is_cuda(), "Inputs must be CUDA");
    TORCH_CHECK(logits.is_contiguous() && labels.is_contiguous(), "Inputs must be contiguous");
    
    int64_t rows = logits.size(0);
    int64_t cols = logits.size(1);

    auto output = torch::empty({rows}, logits.options());

    dim3 grid(rows);
    dim3 block(BLOCK_SIZE);

    poly1_loss_kernel<<<grid, block>>>(
        output.data_ptr<float>(),
        logits.data_ptr<float>(),
        labels.data_ptr<int64_t>(),
        cols,
        epsilon
    );

    if (reduction == "mean") {
        return output.mean();
    } else if (reduction == "sum") {
        return output.sum();
    }

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, epsilon=2.0, reduction='none'):
        super(ModelNew, self).__init__()
        self.epsilon = epsilon
        self.reduction = reduction
        self.op = load_inline(
            name='poly1_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['poly1_loss_cuda_forward'],
            verbose=False,
            extra_cuda_cflags=['-O3']
        )

    def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        return self.op.poly1_loss_cuda_forward(
            logits.contiguous(), 
            labels.contiguous(), 
            self.epsilon, 
            self.reduction
        )