# cosineloss_cuda_ultimate.py
import torch
from torch.utils.cpp_extension import load_inline
from cosineloss_torch import BATCH_SIZE, EMBEDDING_DIM, DIM, MARGIN

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 cosine_forward_cuda(torch::Tensor x1, torch::Tensor x2, torch::Tensor y);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <device_launch_parameters.h>

        #define BLOCK_SIZE 256
        #define VEC_SIZE 4
        #define WARP_SIZE 32
        #define MARGIN_VAL {margin_val}f

        // Fast warp reduction
        __device__ __forceinline__ float warp_reduce_sum(float val) {{
            #pragma unroll
            for (int offset = 16; offset > 0; offset >>= 1) {{
                val += __shfl_down_sync(0xffffffff, val, offset);
            }}
            return val;
        }}

        // Fast block reduction
        __device__ __forceinline__ float block_reduce_sum(float val) {{
            __shared__ float shared[WARP_SIZE];
            
            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;
        }}

        // Single-pass fused kernel - computes everything in one go
        __global__ void cosine_loss_fused_kernel(
            const float* __restrict__ x1,
            const float* __restrict__ x2,
            const float* __restrict__ y,
            float* __restrict__ output_loss,
            const int N_pairs,
            const int D_emb
        ) {{
            const int pair_idx = blockIdx.x;
            
            if (pair_idx >= N_pairs) return;

            // Use float for better performance
            float thread_dot = 0.0f;
            float thread_norm1 = 0.0f;
            float thread_norm2 = 0.0f;

            const int offset = pair_idx * D_emb;
            const int D_vec = D_emb / VEC_SIZE;

            const float4* __restrict__ x1_4 = reinterpret_cast<const float4*>(x1 + offset);
            const float4* __restrict__ x2_4 = reinterpret_cast<const float4*>(x2 + offset);

            // Vectorized computation
            for (int d_vec = threadIdx.x; d_vec < D_vec; d_vec += blockDim.x) {{
                float4 v1 = __ldg(&x1_4[d_vec]);
                float4 v2 = __ldg(&x2_4[d_vec]);
                
                // Dot product with FMA
                thread_dot = fmaf(v1.x, v2.x, thread_dot);
                thread_dot = fmaf(v1.y, v2.y, thread_dot);
                thread_dot = fmaf(v1.z, v2.z, thread_dot);
                thread_dot = fmaf(v1.w, v2.w, thread_dot);

                // Norm squared with FMA
                thread_norm1 = fmaf(v1.x, v1.x, thread_norm1);
                thread_norm1 = fmaf(v1.y, v1.y, thread_norm1);
                thread_norm1 = fmaf(v1.z, v1.z, thread_norm1);
                thread_norm1 = fmaf(v1.w, v1.w, thread_norm1);

                thread_norm2 = fmaf(v2.x, v2.x, thread_norm2);
                thread_norm2 = fmaf(v2.y, v2.y, thread_norm2);
                thread_norm2 = fmaf(v2.z, v2.z, thread_norm2);
                thread_norm2 = fmaf(v2.w, v2.w, thread_norm2);
            }}

            // Block-level reduction
            float dot = block_reduce_sum(thread_dot);
            float norm1_sq = block_reduce_sum(thread_norm1);
            float norm2_sq = block_reduce_sum(thread_norm2);

            // Only thread 0 computes the final loss for this pair
            if (threadIdx.x == 0) {{
                float label = y[pair_idx];
                
                // Compute cosine similarity with safe division
                float norm_prod = sqrtf(norm1_sq * norm2_sq);
                float cosine = (norm_prod > 1e-8f) ? (dot / norm_prod) : 0.0f;
                
                // Compute loss based on label
                float loss;
                if (label > 0.0f) {{
                    loss = 1.0f - cosine;
                }} else {{
                    // Use fmaxf for max(0, cosine - margin)
                    loss = fmaxf(0.0f, cosine - MARGIN_VAL);
                }}
                
                output_loss[pair_idx] = loss;
            }}
        }}

        // Final reduction kernel
        __global__ void final_sum_kernel(
            const float* __restrict__ losses,
            float* __restrict__ output,
            const int n
        ) {{
            float sum = 0.0f;
            
            for (int i = threadIdx.x; i < n; i += WARP_SIZE) {{
                sum += losses[i];
            }}
            
            sum = warp_reduce_sum(sum);
            
            if (threadIdx.x == 0) {{
                output[0] = sum;
            }}
        }}

        torch::Tensor cosine_forward_cuda(torch::Tensor x1, torch::Tensor x2, torch::Tensor y) {{
            TORCH_CHECK(x1.is_cuda() && x2.is_cuda() && y.is_cuda(), 
                       "Inputs must be CUDA tensors");
            
            x1 = x1.contiguous();
            x2 = x2.contiguous();
            y = y.contiguous();

            const int N_pairs = x1.size(0);
            const int D_emb = x1.size(1);
            
            TORCH_CHECK(D_emb % VEC_SIZE == 0, 
                       "Embedding dimension must be divisible by 4");

            // Allocate temp storage for per-pair losses
            auto per_pair_loss = torch::empty({{N_pairs}}, x1.options());
            auto final_result = torch::empty({{1}}, x1.options());

            // Launch fused kernel - one block per pair
            const int block_size = BLOCK_SIZE;
            const int grid_size = N_pairs;

            cosine_loss_fused_kernel<<<grid_size, block_size>>>(
                x1.data_ptr<float>(),
                x2.data_ptr<float>(),
                y.data_ptr<float>(),
                per_pair_loss.data_ptr<float>(),
                N_pairs,
                D_emb
            );

            // Final reduction
            final_sum_kernel<<<1, WARP_SIZE>>>(
                per_pair_loss.data_ptr<float>(),
                final_result.data_ptr<float>(),
                N_pairs
            );

            // Compute mean on GPU
            final_result.div_(N_pairs);

            return final_result;
        }}
        """.format(margin_val=MARGIN)

        self.cos_op = load_inline(
            name="cosine_ultimate_op",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["cosine_forward_cuda"],
            extra_cuda_cflags=[
                "-O3",
                "--use_fast_math",
                "-lineinfo",
            ],
            verbose=True
        )

    def forward(self, x1: torch.Tensor, x2: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return self.cos_op.cosine_forward_cuda(x1, x2, y)