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 ldam_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& targets,
    const torch::Tensor& class_margins,
    float s,
    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;
};

template<typename T>
__device__ __forceinline__ T warp_reduce_max(T val) {
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        val = max(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    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_max(float val) {
    static __shared__ float shared[32];
    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;
    val = warp_reduce_max(val);
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : -FLT_MAX;
    if (wid == 0) val = warp_reduce_max(val);
    return val;
}

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

__global__ void ldam_loss_kernel(
    float* __restrict__ output,
    const float* __restrict__ logits,
    const int64_t* __restrict__ targets,
    const float* __restrict__ class_margins,
    int cols,
    float s) 
{
    int row = blockIdx.x;
    int tid = threadIdx.x;

    const float* row_logits = logits + row * cols;
    int64_t target_idx = targets[row];
    
    // Retrieve margin for the target class of this sample
    // Only used when idx matches target, but let's just read it in thread 0 or lazily.
    // To avoid branch divergence inside loops, we can read it once.
    // But actually, only one thread needs it.
    
    // 1. Pass 1: Find Max (of MODIFIED logits)
    float local_max = -FLT_MAX;
    
    // Grid-Stride Loop
    int i = tid * 4;
    while (i < cols) {
        if (i + 4 <= cols) {
            Float4 vec = reinterpret_cast<const Float4*>(&row_logits[i])[0];
            
            // Unpack to apply logic
            float v0 = vec.x; float v1 = vec.y; float v2 = vec.z; float v3 = vec.w;
            
            // LDAM Logic: if target, sub margin
            if (i == target_idx) v0 -= class_margins[target_idx];
            if (i+1 == target_idx) v1 -= class_margins[target_idx];
            if (i+2 == target_idx) v2 -= class_margins[target_idx];
            if (i+3 == target_idx) v3 -= class_margins[target_idx];
            
            // Scale
            local_max = fmaxf(local_max, v0 * s);
            local_max = fmaxf(local_max, v1 * s);
            local_max = fmaxf(local_max, v2 * s);
            local_max = fmaxf(local_max, v3 * s);
        } else {
            // Tail
            for (int k = 0; k < 4 && (i + k) < cols; ++k) {
                float val = row_logits[i+k];
                if (i + k == target_idx) {
                    val -= class_margins[target_idx];
                }
                val *= s;
                local_max = fmaxf(local_max, val);
            }
        }
        i += blockDim.x * 4;
    }
    
    float global_max = block_reduce_max(local_max);
    
    __shared__ float s_max;
    if (tid == 0) s_max = global_max;
    __syncthreads();
    global_max = s_max;

    // 2. Pass 2: Sum Exp
    float local_sum = 0.0f;
    
    i = tid * 4;
    while (i < cols) {
        if (i + 4 <= cols) {
            Float4 vec = reinterpret_cast<const Float4*>(&row_logits[i])[0];
            float v0 = vec.x; float v1 = vec.y; float v2 = vec.z; float v3 = vec.w;

            if (i == target_idx) v0 -= class_margins[target_idx];
            if (i+1 == target_idx) v1 -= class_margins[target_idx];
            if (i+2 == target_idx) v2 -= class_margins[target_idx];
            if (i+3 == target_idx) v3 -= class_margins[target_idx];
            
            local_sum += __expf(v0 * s - global_max);
            local_sum += __expf(v1 * s - global_max);
            local_sum += __expf(v2 * s - global_max);
            local_sum += __expf(v3 * s - global_max);
        } else {
            for (int k = 0; k < 4 && (i + k) < cols; ++k) {
                float val = row_logits[i+k];
                if (i + k == target_idx) val -= class_margins[target_idx];
                val *= s;
                local_sum += __expf(val - global_max);
            }
        }
        i += blockDim.x * 4;
    }

    float global_sum = block_reduce_sum(local_sum);

    // 3. Final Loss
    if (tid == 0) {
        float target_val = row_logits[target_idx];
        // Modify target
        target_val = (target_val - class_margins[target_idx]) * s;
        
        // CE Loss = -log(e^t / Sum) = log(Sum) - t
        //         = log(sum_exp * e^M) - t
        //         = M + log(sum_exp) - t
        float loss = global_max + __logf(global_sum) - target_val;
        output[row] = loss;
    }
}

torch::Tensor ldam_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& targets,
    const torch::Tensor& class_margins,
    float s,
    std::string reduction) 
{
    TORCH_CHECK(logits.is_cuda() && targets.is_cuda() && class_margins.is_cuda(), "Inputs must be CUDA");
    TORCH_CHECK(logits.is_contiguous() && targets.is_contiguous() && class_margins.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());

    ldam_loss_kernel<<<batch_size, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        logits.data_ptr<float>(),
        targets.data_ptr<int64_t>(),
        class_margins.data_ptr<float>(),
        num_classes,
        s
    );

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

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, cls_num_list, max_m=0.5, s=30.0, reduction='none'):
        super(ModelNew, self).__init__()
        self.s = s
        self.reduction = reduction
        
        # Weights logic is same
        import numpy as np
        m_list = 1.0 / np.sqrt(np.sqrt(cls_num_list))
        m_list = m_list * (max_m / np.max(m_list))
        m_list = torch.FloatTensor(m_list)
        self.register_buffer('class_margins', m_list)
        
        self.op = load_inline(
            name='ldam_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['ldam_loss_cuda_forward'],
            verbose=False,
            extra_cuda_cflags=['-O3']
        )

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        return self.op.ldam_loss_cuda_forward(
            logits.contiguous(), 
            targets.contiguous(), 
            self.class_margins.contiguous(),
            self.s, 
            self.reduction
        )