"""Keep detector proposal TopK as a native TFLite operation, not StableHLO sort."""
import torch
from litert_torch import fx_infra
from litert_torch.backend.lowerings import registry, utils
from litert_torch.backend.experimental.torch_tfl._lowerings import _ir_operation
from litert_converter.mlir import ir

@torch.library.custom_op('vds_static::topk', mutates_args=())
def native_topk(x: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]:
    return torch.topk(x, k, dim=-1, largest=True, sorted=True)

@native_topk.register_fake
def fake_topk(x, k):
    shape = (*x.shape[:-1], k)
    return x.new_empty(shape), x.new_empty(shape, dtype=torch.int64)

@registry.lower(torch.ops.vds_static.topk.default)
def lower_topk(ctx, x, k):
    values, indices = utils.node_meta_to_ir_types(ctx.node)
    i32 = ir.IntegerType.get_signless(32)
    index_type = ir.RankedTensorType.get(indices.shape, i32)
    result = _ir_operation('tfl.topk_v2', results=[values, index_type],
                           operands=[x, utils.splat(k, i32, [])])
    return result[0], _ir_operation('tfl.cast', results=[indices], operands=[result[1]])

def topk(x, k, dim=-1, largest=True, sorted=True):
    dim = dim % x.ndim
    moved = x.movedim(dim, -1)
    values, indices = torch.ops.vds_static.topk(moved if largest else -moved, k)
    if not largest:
        values = -values
    return values.movedim(-1, dim), indices.movedim(-1, dim)

fx_infra.decomp.update_pre_convert_decomp({torch.ops.aten.topk.default: topk})
fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.topk.default: topk})