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

cpp_source = """
#include <torch/extension.h>
#include <string>

torch::Tensor hard_bootstrap_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& targets,
    float beta,
    std::string reduction);
"""

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

#define BLOCK_SIZE 256
#define WARP_SIZE 32

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

struct MaxPair {
    float val;
    int idx;
};

__device__ __forceinline__ MaxPair reduce_max_pair_op(MaxPair a, MaxPair b) {
    if (a.val > b.val) return a;
    return b;
}

__device__ __forceinline__ MaxPair warp_reduce_argmax(MaxPair val) {
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        float other_val = __shfl_down_sync(0xffffffff, val.val, offset);
        int other_idx = __shfl_down_sync(0xffffffff, val.idx, offset);
        
        if (other_val > val.val) {
            val.val = other_val;
            val.idx = other_idx;
        }
    }
    return val;
}

__device__ __forceinline__ MaxPair block_reduce_argmax(MaxPair val) {
    static __shared__ MaxPair shared[32];
    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;

    val = warp_reduce_argmax(val);
    if (lane == 0) shared[wid] = val;
    __syncthreads();

    // Last warp reduction
    if (wid == 0) {
        val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : MaxPair{-FLT_MAX, -1};
        val = warp_reduce_argmax(val);
    }
    return val;
}

template<typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__device__ __forceinline__ float block_reduce_sum(float val) {
    static __shared__ float shared_sum[32];
    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;
    val = warp_reduce_sum(val);
    if (lane == 0) shared_sum[wid] = val;
    __syncthreads();
    
    if (wid == 0) {
        val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared_sum[lane] : 0.0f;
        val = warp_reduce_sum(val);
    }
    return val;
}

__global__ void hard_bootstrap_kernel(
    float* __restrict__ output,
    const float* __restrict__ logits,
    const int64_t* __restrict__ targets,
    int cols,
    float beta) 
{
    // Cache row in Shared Memory
    extern __shared__ float s_row[];

    int row_idx = blockIdx.x;
    int tid = threadIdx.x;

    const float* row_logits = logits + row_idx * cols;
    
    // 1. Load Data
    int i = tid * 4;
    while (i < cols) {
        if (i + 4 <= cols) {
            *(Float4*)&s_row[i] = *(reinterpret_cast<const Float4*>(&row_logits[i]));
        } else {
            for (int k = 0; k < 4 && i + k < cols; ++k) {
                s_row[i + k] = row_logits[i + k];
            }
        }
        i += blockDim.x * 4;
    }
    __syncthreads();

    // 2. Find Max & ArgMax
    MaxPair local_max = {-FLT_MAX, -1};
    for (int k = tid; k < cols; k += blockDim.x) {
        float v = s_row[k];
        if (v > local_max.val) {
            local_max.val = v;
            local_max.idx = k;
        }
    }
    MaxPair global_max_pair = block_reduce_argmax(local_max);
    
    __shared__ float s_max_val;
    __shared__ int s_max_idx;
    if (tid == 0) {
        s_max_val = global_max_pair.val;
        s_max_idx = global_max_pair.idx;
    }
    __syncthreads();
    
    float M = s_max_val;
    // Note: s_max_idx is used by Thread 0 later

    // 3. Compute Sum Exp
    float local_sum = 0.0f;
    for (int k = tid; k < cols; k += blockDim.x) {
        local_sum += __expf(s_row[k] - M);
    }
    float S = block_reduce_sum(local_sum);

    // 4. Compute Loss
    if (tid == 0) {
        float log_S = __logf(S);
        int64_t target_idx = targets[row_idx];
        // Ensure index is valid (safety check)
        int argmax_idx = s_max_idx; 
        
        float val_target = s_row[target_idx];
        float val_argmax = s_max_val; // Argmax value is M
        
        // L_ce = -log(p_t) = log(S) + M - x_t
        // L_boot = -log(p_max) = log(S) + M - x_max
        // L = beta * L_ce + (1-beta) * L_boot
        //   = log(S) + M - (beta * x_t + (1-beta) * x_max)
        
        float weighted_logit = beta * val_target + (1.0f - beta) * val_argmax;
        float loss = log_S + M - weighted_logit;
        
        output[row_idx] = loss;
    }
}

torch::Tensor hard_bootstrap_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& targets,
    float beta,
    std::string reduction) 
{
    TORCH_CHECK(logits.is_cuda() && targets.is_cuda(), "Inputs must be CUDA");
    TORCH_CHECK(logits.is_contiguous() && targets.is_contiguous(), "Inputs must be contiguous");
    
    int batch_size = logits.size(0);
    int num_classes = logits.size(1);

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

    size_t smem_size = num_classes * sizeof(float);
    
    hard_bootstrap_kernel<<<batch_size, BLOCK_SIZE, smem_size>>>(
        output.data_ptr<float>(),
        logits.data_ptr<float>(),
        targets.data_ptr<int64_t>(),
        num_classes,
        beta
    );

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

    return output;
}
"""

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

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        return self.op.hard_bootstrap_loss_cuda_forward(
            logits.contiguous(), 
            targets.contiguous(), 
            self.beta, 
            self.reduction
        )