# marginrankingloss_cuda.py
import torch
from torch.utils.cpp_extension import load_inline
from marginrankingloss_torch import BATCH_SIZE, FEATURE_DIM, MARGIN 



N_ELEMENTS = BATCH_SIZE * FEATURE_DIM

class ModelNew(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor mrl_forward_cuda(torch::Tensor x1, torch::Tensor x2, torch::Tensor target, float margin_val);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <cmath>
        #include <float.h>

        #define BLOCK_SIZE 256
        #define MARGIN_VAL {margin_val}f 

        __device__ __forceinline__ float warp_reduce_sum(float val) {{ 
            for (int offset = BLOCK_SIZE / 2; offset > 0; offset /= 2) {{ 
                val += __shfl_down_sync(0xffffffff, val, offset);
            }} 
            return val;
        }}


        __global__ void mrl_fused_kernel(
            const float* __restrict__ x1,
            const float* __restrict__ x2,
            const float* __restrict__ target,
            float* __restrict__ loss_sum_out, 
            int n_elements,
            float margin_val 
        ) {{ 
            __shared__ float s_data[BLOCK_SIZE];
            
            // Grid-Stride Loop
            float thread_loss = 0.0f;
            int grid_stride = gridDim.x * blockDim.x;

            for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
                idx < n_elements;
                idx += grid_stride)
            {{ 
                float val1 = x1[idx];
                float val2 = x2[idx];
                float y = target[idx];

                // Loss = max(0, -y * (x1 - x2) + margin)
                
                float diff = val1 - val2;
                float term = -y * diff + MARGIN_VAL; 
                
                float loss_val = fmaxf(0.0f, term);
                
                thread_loss += loss_val;
            }} 
            
            s_data[threadIdx.x] = thread_loss;
            __syncthreads();

            for (int offset = BLOCK_SIZE / 2; offset > 0; offset >>= 1) {{ 
                if (threadIdx.x < offset) {{ 
                    s_data[threadIdx.x] += s_data[threadIdx.x + offset];
                }} 
                __syncthreads();
            }} 

            if (threadIdx.x == 0) {{ 
                loss_sum_out[blockIdx.x] = s_data[0]; 
            }} 
        }} 

        torch::Tensor mrl_forward_cuda(torch::Tensor x1, torch::Tensor x2, torch::Tensor target, float margin_val) {{ 
            TORCH_CHECK(x1.is_cuda(), "Input must be a CUDA tensor");
            x1 = x1.contiguous();
            x2 = x2.contiguous();
            target = target.contiguous();

            int n_elements = x1.numel();
            
            const int block_size = BLOCK_SIZE;
            const int grid_size = std::max(1, (n_elements + block_size - 1) / block_size); 

            auto block_loss_sums = torch::empty({{grid_size}}, x1.options()); 

            mrl_fused_kernel<<<grid_size, block_size>>>(
                x1.data_ptr<float>(),
                x2.data_ptr<float>(),
                target.data_ptr<float>(),
                block_loss_sums.data_ptr<float>(),
                n_elements,
                margin_val
            );
            

            return block_loss_sums.sum() / n_elements;
        }} 
        """.format(margin_val=MARGIN) 
        
        self.mrl_op = load_inline(
            name="mrl_fused_op_safest",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["mrl_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=True
        )

    def forward(self, x1: torch.Tensor, x2: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        return self.mrl_op.mrl_forward_cuda(x1, x2, target, MARGIN)