from typing import Callable, Dict, List, Set

import torch

import torch.fx as fx

import torch.utils._pytree as pytree

from torch import Tensor

from torch.distributed._tensor import DeviceMesh, Replicate, Shard
from torch.distributed._tensor.ops.view_ops import (
    DimSpec,
    InputDim,
    ops as view_op_rules,
)
from torch.distributed._tensor.placement_types import _Partial, DTensorSpec

aten = torch.ops.aten


class BatchDimAnalyzer:
    """
    This class is used to analyze the batch dimension of each tensor/node in the
    graph. We need to know the batch dimension of each tensor/node so that we know
    exactly the sharding layout of intermediate tensors.

    We possibly should evaluate using symbolic shapes to track the batch dimension.
    We can experiment it later with dynamo integration (as dynamo have mark_dynamic
    API which allows marking batch dimension only) or try to use FakeTensorMode to
    mark the batch dimension. For now, let's just use the batch dimension of the first
    input tensor as the hint to track the batch dimension of all tensors/nodes in
    the graph.
    """

    def __init__(self, batch_dim: int = 0) -> None:
        self.batch_dim = batch_dim

        self.batch_dim_map: Dict[fx.Node, int] = {}
        # batch dim size is used to track the batch dim size of the input tensor
        self.batch_dim_size = -1

        self.dim_rule_map: Dict[torch._ops.OpOverload, Callable[..., torch.Tensor]] = {
            aten.squeeze.default: torch.squeeze,
            aten.squeeze.dim: torch.squeeze,
            aten.view.default: Tensor.view,
            aten.reshape.default: torch.reshape,
            aten._unsafe_view.default: Tensor.view,
            aten.unsqueeze.default: torch.unsqueeze,
            aten.expand.default: Tensor.expand,
            aten.permute.default: torch.permute,
            aten.repeat.default: Tensor.repeat,
            aten.transpose.int: torch.transpose,
        }

    def init_batch_dim_size(self, batch_dim_size: int) -> None:
        """
        initialize batch dim size base on the first input batch size
        """
        if self.batch_dim_size != -1 and self.batch_dim_size != batch_dim_size:
            raise RuntimeError(
                f"batch dim size is already initialized! "
                f"Found new batch size: {batch_dim_size} not "
                f"matching existing batch dim size: {self.batch_dim_size}!"
            )
        self.batch_dim_size = batch_dim_size

    def set_batch_dim(self, node: fx.Node, batch_dim: int) -> None:
        self.batch_dim_map[node] = batch_dim

    def get_batch_dim(self, node: fx.Node) -> int:
        if node not in self.batch_dim_map:
            raise RuntimeError(f"batch dim analysis failed on node: {node}!")
        return self.batch_dim_map[node]

    def compute_batch_dim(self, node: fx.Node, full_reduction=False) -> int:
        """
        compute the batch dimension for the `node`
        """
        assert self.batch_dim_size != -1, "batch dim size is not initialized!"

        if node in self.batch_dim_map:
            # if batch dim already computed, simply return it
            return self.batch_dim_map[node]

        if node.target in self.dim_rule_map:
            view_op_rule = view_op_rules[self.dim_rule_map[node.target]]  # type: ignore[index]
            args_val = pytree.tree_map_only(fx.Node, lambda n: n.meta["val"], node.args)
            kwargs_val = pytree.tree_map_only(
                fx.Node, lambda n: n.meta["val"], node.kwargs
            )
            output_dim_rules = view_op_rule.dim_map(*args_val, **kwargs_val)

            def collect_input_dim(cmd: DimSpec, input_dims: Set[int]):
                if isinstance(cmd, InputDim):
                    input_dims.add(cmd.input_dim)
                for inp in cmd.inputs():
                    collect_input_dim(inp, input_dims)

            output_dim_to_input_dims: List[Set[int]] = []
            for inp in output_dim_rules:
                input_dims: Set[int] = set()
                collect_input_dim(inp, input_dims=input_dims)
                output_dim_to_input_dims.append(input_dims)

            operand = node.all_input_nodes[0]
            operand_batch_dim = self.get_batch_dim(operand)
            for output_dim, input_dims in enumerate(output_dim_to_input_dims):
                if operand_batch_dim in input_dims:
                    self.set_batch_dim(node, output_dim)
                    # update batch dim size before return
                    # this is because batch dim size might change during the middle
                    self.batch_dim_size = node.meta["val"].shape[output_dim]
                    return output_dim

        # if there's no hints from the output_dim_rules, we infer from output
        # shape to see if there's batch dim, and shard correspondingly
        node_val = node.meta["val"]
        if isinstance(node_val, (list, tuple)):
            shapes = [val.shape for val in node_val]
        else:
            shapes = [node_val.shape]

        # for reduction op that reduces over the sharded batch dim
        # we don't generate partial, but rather, we generate shard
        # This is because the intention of data parallel is to never
        # do full reduction across batch dimension, it would still
        # keep the reduction activation as sharded.
        full_reduction = False
        # loop through the dim size to find the output batch dim
        for shape in shapes:
            if len(shape) == 0:
                full_reduction = True

            for i, dim_size in enumerate(shape):
                if dim_size == self.batch_dim_size:
                    self.set_batch_dim(node, i)
                    return i

        operands = node.all_input_nodes
        if not operands:
            # if there's no operands, it must be factory ops and it's a tensor
            # generated for computation and should be marked as replicated
            self.set_batch_dim(node, -1)
            # -1 means replicated
            return -1
        else:
            # if there's operand we see the operand have batch dim, if operand
            # have batch dim but output does not, it's either a full reduction,
            # where we should stay sharded, or it's a reduction on batch dim only
            # where we should produce partial
            operand_batch_dim = -1
            for operand in operands:
                if operand in self.batch_dim_map:
                    operand_batch_dim = self.get_batch_dim(operand)
            # self.get_batch_dim(operands[0])
            if operand_batch_dim < 0:
                # if operand does not have batch dim, we also don't have batch dim
                self.set_batch_dim(node, operand_batch_dim)
                return operand_batch_dim
            elif full_reduction:
                self.set_batch_dim(node, operand_batch_dim)
                return operand_batch_dim
            else:
                # if operand have batch dim but output does not, it should
                # produce partial, we use -2 to indicate partial
                self.set_batch_dim(node, -2)
                return -2

    def compute_act_spec(self, node: fx.Node, mesh: DeviceMesh) -> DTensorSpec:
        """
        This function first compute the batch dimension for the current node,
        then generate the sharding spec that shards on the batch dimension.
        """
        node_batch_dim = self.compute_batch_dim(node)
        if node_batch_dim == -1:
            # indicate this activation is replicated
            act_spec = DTensorSpec(mesh=mesh, placements=(Replicate(),))
        elif node_batch_dim == -2:
            # indicate this activation is partial
            act_spec = DTensorSpec(mesh=mesh, placements=(_Partial(),))
        else:
            # indicate this activation is Shard
            act_spec = DTensorSpec(mesh=mesh, placements=(Shard(node_batch_dim),))

        return act_spec