# Copyright (c) Meta Platforms, Inc. and affiliates

import torch
from torch import SymInt

from ..device_mesh import DeviceMesh


# Register custom operator
torch.library.define(
    "device_mesh::_runtime_compute_coordinate_on_dim",
    "(Tensor full_mesh, int index) -> SymInt",
    tags=torch.Tag.pt2_compliant_tag,
)


@torch.library.register_fake("device_mesh::_runtime_compute_coordinate_on_dim")
def _runtime_compute_coordinate_on_dim_fake(
    full_mesh: torch.Tensor, index: int
) -> SymInt:
    from torch.fx.experimental.symbolic_shapes import _constrain_range_for_size

    ctx = torch._custom_op.impl.get_ctx()
    shape_env = ctx._shape_env

    # Bypass allow_dynamic_output_shape_ops check by directly creating the symint.
    # This is intentional - the coordinate is always valid and bounded.
    sz = shape_env.create_unbacked_symint()

    # Apply size constraints - coordinate is bounded by mesh size on the given dimension.
    # The full_mesh tensor has an extra batch dimension at the front, so the actual
    # mesh dimensions start at index 1. mesh.size(index) = full_mesh.size(index + 1)
    mesh_size = full_mesh.size(index + 1)
    _constrain_range_for_size(
        sz, min=0, max=mesh_size - 1 if isinstance(mesh_size, int) else None
    )

    try:
        # Check if we're currently tracing in dynamo (as opposed to AOT or export).
        in_dynamo = torch._dynamo.symbolic_convert.InstructionTranslator.current_tx()
    except AttributeError:
        in_dynamo = False

    if in_dynamo:
        # During dynamo tracing, distributed ops are treated as atomic - so the
        # rank SymInt may be computed but not traced into the graph (e.g., it
        # affects tensor values but not shapes). Mark it as ignorable here;
        # when we decompose these ops later (after dynamo), we'll create fresh
        # SymInts that do get traced.
        shape_env.ignorable_fresh_unbacked_symbols.append(sz.node._expr)

    return sz


@torch.library.impl(
    "device_mesh::_runtime_compute_coordinate_on_dim", "CompositeExplicitAutograd"
)
def _runtime_compute_coordinate_on_dim_impl(full_mesh: torch.Tensor, index: int) -> int:
    rank = torch.distributed.get_rank()
    mesh = DeviceMesh._get_mesh_tensor_from_full_mesh(full_mesh)
    mesh_coords = DeviceMesh._compute_coordinates_from_mesh(mesh, rank)
    if mesh_coords is None:
        raise AssertionError
    return mesh_coords[index]


def _get_flattened_submesh_impl(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
    from torch.distributed.tensor._redistribute import (
        _get_flattened_mesh_by_layout_impl,
    )

    result = _get_flattened_mesh_by_layout_impl(mesh, tuple(mesh_dims))
    if result is None:
        raise ValueError(f"No flattened mesh found for mesh_dims={mesh_dims} on {mesh}")
    return result


@torch.library.custom_op("device_mesh::_get_flattened_submesh", mutates_args=())
def _get_flattened_submesh(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
    return _get_flattened_submesh_impl(mesh, mesh_dims)


@_get_flattened_submesh.register_fake
def _get_flattened_submesh_fake(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
    return _get_flattened_submesh_impl(mesh, mesh_dims)


def _get_submesh_impl(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
    all_dim_names = mesh._mesh_dim_names
    if all_dim_names is None:
        raise ValueError(f"Cannot slice mesh without dim names: {mesh}")
    dim_names = tuple(all_dim_names[i] for i in mesh_dims)
    if len(dim_names) == 1:
        return mesh[dim_names[0]]
    return mesh[dim_names]


@torch.library.custom_op("device_mesh::_get_submesh", mutates_args=())
def _get_submesh(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
    return _get_submesh_impl(mesh, mesh_dims)


@_get_submesh.register_fake
def _get_submesh_fake(mesh: DeviceMesh, mesh_dims: list[int]) -> DeviceMesh:
    return _get_submesh_impl(mesh, mesh_dims)
