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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor max_rank_forward(torch::Tensor x);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void max_rank4_kernel(const float* __restrict__ x, float* __restrict__ out) {
    int i = threadIdx.x + blockIdx.x * blockDim.x;
    if (i >= 224) return;

    // x layout is contiguous (R=4, B=8, C=28), so each rank plane has 224 elements.
    float best = x[i];
    int idx = 0;

    float v1 = x[224 + i];
    if (v1 > best) { best = v1; idx = 1; }
    float v2 = x[448 + i];
    if (v2 > best) { best = v2; idx = 2; }
    float v3 = x[672 + i];
    if (v3 > best) { best = v3; idx = 3; }

    int o = i * 2;
    out[o] = best;
    out[o + 1] = (float)idx;
}

torch::Tensor max_rank_forward(torch::Tensor x) {
    auto out = torch::empty({8, 28, 2}, x.options());
    const int threads = 256;
    max_rank4_kernel<<<1, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
"""

_ext = load_inline(
    name="max_rank4_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["max_rank_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.max_rank_forward(x)