"""
Over virtual ranks R, for each (B, C) position take the maximum value; on ties choose the smallest rank index.

Output (B, C, 2): [:,:,0] = best value, [:,:,1] = winning rank as float32.
"""

import torch
import torch.nn as nn

from kernelbench.distributed_collectives import default_device

R = 4
B = 8
C = 28


class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (R, B, C)
        best_val = x[0].clone()
        best_idx = torch.zeros(B, C, dtype=torch.long, device=x.device)
        for r in range(1, R):
            xv = x[r]
            improved = xv > best_val
            tied = (xv == best_val) & ~improved
            best_idx = torch.where(
                improved,
                torch.full_like(best_idx, r),
                torch.where(
                    tied,
                    torch.minimum(best_idx, torch.full_like(best_idx, r)),
                    best_idx,
                ),
            )
            best_val = torch.where(improved, xv, best_val)
        stacked = torch.stack([best_val, best_idx.to(best_val.dtype)], dim=-1)
        return stacked


def get_inputs():
    dev = default_device()
    p = popcorn_pri
    return [torch.randn(R, p.jitter_int(B), p.jitter_int(C), device=dev)]


def get_init_inputs():
    return []

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

# Compile the CUDA kernel with vectorized loads
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

// Vectorized version: each thread handles 4 consecutive (B, C) positions
__global__ void __launch_bounds__(128) max_with_tiebreak_vector_kernel(
    const float* __restrict__ x,  // shape (R, B, C), contiguous
    float* __restrict__ out,       // shape (B, C, 2), contiguous
    int R, int B, int C
) {
    int bc_idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    int total = B * C;
    
    // Use float4 for vectorized loads
    int stride = B * C;
    
    // Process 4 elements per thread
    #pragma unroll
    for (int i = 0; i < 4; i++) {
        int idx = bc_idx + i;
        if (idx >= total) return;
        
        float best_val = __ldg(x + idx);
        int best_idx = 0;
        
        #pragma unroll 4
        for (int r = 1; r < R; r++) {
            float val = __ldg(x + r * stride + idx);
            if (val > best_val) {
                best_val = val;
                best_idx = r;
            } else if (val == best_val) {
                if (r < best_idx) {
                    best_idx = r;
                }
            }
        }
        
        // Write output
        out[idx * 2] = best_val;
        out[idx * 2 + 1] = (float)best_idx;
    }
}

void launch_max_with_tiebreak(const torch::Tensor& x, torch::Tensor& out) {
    int R = x.size(0);
    int B = x.size(1);
    int C = x.size(2);
    
    int total = B * C;
    int threads = 128;
    int blocks = (total / 4 + threads - 1) / threads + 1;  // +1 for remainder
    
    max_with_tiebreak_vector_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        out.data_ptr<float>(),
        R, B, C
    );
}
"""

cpp_source = """
void launch_max_with_tiebreak(const torch::Tensor& x, torch::Tensor& out);
"""

# Load the extension
max_tiebreak_ext = load_inline(
    name="max_tiebreak_ext_v3",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["launch_max_with_tiebreak"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.ext = max_tiebreak_ext
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (R, B, C)
        R, B, C = x.shape
        out = torch.empty(B, C, 2, dtype=torch.float32, device=x.device)
        self.ext.launch_max_with_tiebreak(x, out)
        return out