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 sparsemax_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& targets,
    std::string reduction);
"""

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

#define BLOCK_SIZE 1024
#define NUM_ELEM 2048

__device__ inline void swap(float& a, float& b) {
    float tmp = a; a = b; b = tmp;
}

__global__ void sparsemax_loss_kernel(
    float* __restrict__ output,
    const float* __restrict__ logits,
    const int64_t* __restrict__ targets,
    int cols) 
{
    // 1. Load Logits into Shared Memory
    __shared__ float s_val[NUM_ELEM];
    
    int row_idx = blockIdx.x;
    int tid = threadIdx.x;
    
    const float* row_logits = logits + row_idx * cols;
    
    // Retrieve Target Logit early (before we mess up indices or sort)
    int64_t target_idx = targets[row_idx];
    float z_y = row_logits[target_idx]; // Global read

    // Load 2 elements per thread
    int idx1 = tid;
    int idx2 = tid + BLOCK_SIZE;
    
    s_val[idx1] = row_logits[idx1];
    s_val[idx2] = row_logits[idx2];

    __syncthreads();

    // 2. Bitonic Sort (Descending)
    for (int size = 2; size <= NUM_ELEM; size <<= 1) {
        // Bitonic Merge
        // Descending order: means we want largest first.
        // dir = ( (tid & (size / 2)) == 0 ) check is for alternating up/down
        // But for the final full merge, we want one direction.
        // Bitonic sort produces a monotonic sequence only at the very end.
        
        for (int stride = size / 2; stride > 0; stride >>= 1) {
            __syncthreads();
            // Emulate 2048 threads with 1024 threads loop
            // Effective thread ID mapping for bitonic network
            
            // Algorithm:
            // For each pair (pos, pos+stride)
            // Logic for "Standard" Bitonic Sort implementation (iterative):
            // There are NUM_ELEM / 2 comparators. We have 1024 threads. Perfect match.
            
            int pos = 2 * tid - (tid & (stride - 1));
            int partner = pos + stride;
            
            if (partner < NUM_ELEM) {
                 float a = s_val[pos];
                 float b = s_val[partner];
                 
                 // Direction logic:
                 // Full sort descending: goal is for final stage to be descending
                 // The XOR trick determines direction for sub-blocks
                 bool sort_descending = ((pos & size) == 0);
                 
                 // If size == NUM_ELEM, we force direction to be Descending (or Ascending depending on what we want)
                 // Actually for standard bitonic sort, the direction flag flips.
                 // To get a fully Descending array:
                 // We essentially run standard sort but invert compare.
                 
                 // Let's keep it simple: Standard Bitonic creates Ascending.
                 // To get Descending, we swap logic.
                 if (sort_descending) {
                     if (a < b) { s_val[pos] = b; s_val[partner] = a; }
                 } else {
                     if (a > b) { s_val[pos] = b; s_val[partner] = a; }
                 }
            }
        }
    }
    __syncthreads();
    
    // Now s_val is Sorted Descending: z_(1) >= z_(2) ... >= z_(K)

    // 3. Parallel Prefix Sum (Scan) - Inclusive
    // We need cumsum to check condition: 1 + k * z_k > cumsum_k
    // Use Hillis-Steele double buffering
    __shared__ float s_sum[2][NUM_ELEM];
    
    // Init scan buffer
    s_sum[0][idx1] = s_val[idx1];
    s_sum[0][idx2] = s_val[idx2];
    __syncthreads();

    int in_buf = 0;
    int out_buf = 1;

    for (int stride = 1; stride < NUM_ELEM; stride <<= 1) {
        __syncthreads(); // barrier between steps
        
        // Process idx1
        if (idx1 >= stride) 
            s_sum[out_buf][idx1] = s_sum[in_buf][idx1] + s_sum[in_buf][idx1 - stride];
        else 
            s_sum[out_buf][idx1] = s_sum[in_buf][idx1];

        // Process idx2
        if (idx2 >= stride)
            s_sum[out_buf][idx2] = s_sum[in_buf][idx2] + s_sum[in_buf][idx2 - stride];
        else
            s_sum[out_buf][idx2] = s_sum[in_buf][idx2];

        // Swap
        in_buf = 1 - in_buf;
        out_buf = 1 - out_buf;
    }
    __syncthreads();
    // Result is in in_buf

    // 4. Find Threshold k
    // Condition: 1 + k * z_k > cumsum_k
    // k is 1-based index (1..C). Array is 0-based (0..C-1).
    // So for index i: k = i + 1.
    // Cond: 1 + (i+1) * s_val[i] > s_sum[in_buf][i]
    
    // We need to find the LARGEST i satisfying this.
    // Since z is sorted, this property is monotonic.
    // We can simply count how many elements satisfy this.
    
    int satisfy1 = (1.0f + (float)(idx1 + 1) * s_val[idx1] > s_sum[in_buf][idx1]) ? 1 : 0;
    int satisfy2 = (1.0f + (float)(idx2 + 1) * s_val[idx2] > s_sum[in_buf][idx2]) ? 1 : 0;
       
    // Let's put satisfy counts into s_sum[0] and reduce
    s_sum[0][idx1] = (float)satisfy1;
    s_sum[0][idx2] = (float)satisfy2;
    __syncthreads();

    // Tree reduction for K
    for (int s = NUM_ELEM / 2; s > 0; s >>= 1) {
        if (tid < s) {
            // Each thread sums 2 nodes, but stride handling needs care for > BLOCK_SIZE
            // Our threads cover 0..1023. Total 2048.
            // Standard reduction:
            // Iter 1: s=1024. tid 0..1023. Add [tid] and [tid+1024].
            // Iter 2: s=512. tid 0..511. Add [tid] and [tid+512].
            
            // Note: Initial mapping was s_sum[idx1] and s_sum[idx2] where idx2 = idx1 + 1024.
            // So step 1 is just:
            s_sum[0][tid] += s_sum[0][tid + s]; 
        }
        __syncthreads();
    }

    // Now s_sum[0][0] holds k(z)
    __shared__ float k_z_val;
    __shared__ float tau;
    
    if (tid == 0) {
        k_z_val = s_sum[0][0];
        // tau = (cumsum[k-1] - 1) / k
        int k_idx = (int)k_z_val - 1;
        // Retrieve cumsum from buffer. buffer index is in_buf
        float cumsum_val = s_sum[in_buf][k_idx];
        tau = (cumsum_val - 1.0f) / k_z_val;
    }
    __syncthreads();

    // 5. Calculate Loss
    // L = 0.5 * sum_{j in S} (z_j^2 - tau^2) + 0.5 - z_y
    // S is indices 0 to k-1
    
    // Each thread calculates z^2 - tau^2 for its elements IF they are in support
    float local_loss_part = 0.0f;
    float t = tau;
    int k_limit = (int)k_z_val;
    
    if (idx1 < k_limit) {
        float z = s_val[idx1];
        local_loss_part += (z * z - t * t);
    }
    if (idx2 < k_limit) {
        float z = s_val[idx2];
        local_loss_part += (z * z - t * t);
    }
    
    // Reduce loss parts
    s_sum[0][idx1] = local_loss_part; // Reusing buffer 0
    s_sum[0][idx2] = 0.0f; // Clear second slot (since reduction below assumes sum of tid and tid+s)
    // Actually, better: store local sum in s_sum[0][tid] = local_loss_part (which includes idx1 and idx2)
    __syncthreads();
    
    s_sum[0][tid] = local_loss_part;
    __syncthreads();

    // Reduction
    for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            s_sum[0][tid] += s_sum[0][tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        float support_term = s_sum[0][0];
        output[row_idx] = 0.5f * support_term + 0.5f - z_y;
    }
}

torch::Tensor sparsemax_loss_cuda_forward(
    const torch::Tensor& logits, 
    const torch::Tensor& targets,
    std::string reduction) 
{
    TORCH_CHECK(logits.is_cuda() && targets.is_cuda(), "Inputs must be CUDA");
    TORCH_CHECK(logits.is_contiguous(), "Logits must be contiguous");
    TORCH_CHECK(logits.size(1) == NUM_ELEM, "Kernel optimized for 2048 classes");
    
    int batch_size = logits.size(0);
    auto output = torch::empty({batch_size}, logits.options());

    sparsemax_loss_kernel<<<batch_size, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        logits.data_ptr<float>(),
        targets.data_ptr<int64_t>(),
        NUM_ELEM
    );

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

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

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