"""
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 []

"""
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 torch.utils.cpp_extension import load_inline
import os

R = 4
B = 8
C = 28

# Compile the CUDA kernel
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void __launch_bounds__(256) max_over_ranks_kernel(
    const float* __restrict__ x,
    float* __restrict__ output,
    int B, int C) {
    
    // Each thread handles one (b, c) position
    int bc_idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (bc_idx >= B * C) return;
    
    int b = bc_idx / C;
    int c = bc_idx % C;
    
    int stride = B * C;
    int base_offset = b * C + c;
    
    // Load all R=4 values using __ldg for read-only cache
    float v0 = __ldg(&x[base_offset]);
    float v1 = __ldg(&x[base_offset + stride]);
    float v2 = __ldg(&x[base_offset + stride * 2]);
    float v3 = __ldg(&x[base_offset + stride * 3]);
    
    // Find max with smallest index on ties
    float best_val = v0;
    int best_idx = 0;
    
    if (v1 > best_val) { best_val = v1; best_idx = 1; }
    if (v2 > best_val) { best_val = v2; best_idx = 2; }
    if (v3 > best_val) { best_val = v3; best_idx = 3; }
    
    // Output is (B, C, 2)
    int out_base = bc_idx * 2;
    output[out_base + 0] = best_val;
    output[out_base + 1] = (float)best_idx;
}

torch::Tensor max_over_ranks_cuda(torch::Tensor x) {
    TORCH_CHECK(x.dim() == 3, "Input must be 3D");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "Input must be float");
    TORCH_CHECK(x.size(0) == 4, "R dimension must be 4");
    
    int B = x.size(1);
    int C = x.size(2);
    
    auto output = torch::empty({B, C, 2}, x.options());
    
    int total_threads = B * C;
    int threads_per_block = 256;
    int num_blocks = (total_threads + threads_per_block - 1) / threads_per_block;
    
    max_over_ranks_kernel<<<num_blocks, threads_per_block>>>(
        x.data_ptr<float>(),
        output.data_ptr<float>(),
        B, C
    );
    
    return output;
}
"""

cpp_source = """
torch::Tensor max_over_ranks_cuda(torch::Tensor x);
"""

module = load_inline(
    name='max_over_ranks_ext_v4',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['max_over_ranks_cuda'],
    extra_cuda_cflags=['-O3', '-use_fast_math'],
    verbose=True
)


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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return module.max_over_ranks_cuda(x)