"""Sharding operators for tensor parallelism."""

import dataclasses
from contextlib import contextmanager
from typing import Any, Dict, List, Optional  # noqa: UP035

from tvm import te, tirx, topi
from tvm.relax.frontend import nn


@dataclasses.dataclass
class ShardSingleDim:
    """
    Shard a tensor by a single dimension.


    Parameters
    ----------
    name : str
        The name of the shard func

    dim : int
        The dimension to shard

    segs : Optional[List[int]]
        The length of segments along `dim`. Default to None. If specified,
        shard a tensor by its "segmented" dimension, where each segment has a different length
        and sharded evenly on each worker.

    """

    name: str
    dim: int
    segs: Optional[List[int]] = None  # noqa: UP006

    def gen_tir(self, shards: int, weight: nn.Tensor) -> tirx.PrimFunc:
        """Generate a TIR function that shards the weight tensor by its rows."""
        shape = weight.shape
        segs = self.segs or [shape[self.dim]]
        assert sum(segs) == shape[self.dim]
        # NOTE: we use int64 to prevent int32 overflow
        shape = [tirx.IntImm("int64", v) for v in shape]
        segs = [tirx.IntImm("int64", v) for v in segs]
        w = te.placeholder(
            [tirx.IntImm("int64", v) for v in self._compute_in_shape(shards, weight)],
            weight.dtype,
            name="w",
        )
        ws: List[te.Tensor] = []  # noqa: UP006
        offset = 0
        for idx, sub_seg in enumerate(segs):
            ws.append(
                topi.transpose(
                    topi.reshape(
                        te.compute(
                            (
                                *shape[: self.dim],
                                sub_seg * shards,
                                *shape[self.dim + 1 :],
                            ),
                            lambda *idx: w[
                                (
                                    *idx[: self.dim],
                                    idx[self.dim] + offset,
                                    *idx[self.dim + 1 :],
                                )
                            ],
                            name=f"w_{idx}",
                        ),
                        (
                            *shape[: self.dim],
                            tirx.IntImm("int64", shards),
                            sub_seg,
                            *shape[self.dim + 1 :],
                        ),
                    ),
                    [self.dim, *range(self.dim), *range(self.dim + 1, len(shape) + 1)],
                )
            )
            offset += sub_seg * shards
        o = topi.concatenate(ws, axis=1 + self.dim)
        func = te.create_prim_func([w, o])
        return func

    def gen_shard_info(self, shards: int, weight: nn.Tensor) -> Dict[str, Any]:  # noqa: UP006
        """Generate shard info for this sharding strategy."""
        return {
            "func_name": self.name,
            "in_shape": self._compute_in_shape(shards, weight),
            "out_shape": (shards, *weight.shape),
            "out_dtype": str(weight.dtype),
        }

    def _compute_in_shape(self, shards: int, weight: nn.Tensor) -> List[int]:  # noqa: UP006
        """Compute the weight shape before sharding."""
        shape = weight.shape
        return [*shape[: self.dim], shape[self.dim] * shards, *shape[self.dim + 1 :]]


@contextmanager
def shard_bias(linear: nn.Linear, tensor_parallel_shards: int):
    """
    A context manager to shard the bias of a linear into `tensor_parallel_shards` shards.


    Parameters
    ----------
    linear : nn.Linear
        The linear layer whose bias would be sharded.

    tensor_parallel_shards : int
        The number of shards.
    """
    original_bias = linear.bias
    if tensor_parallel_shards > 1:
        linear.bias = linear.bias / tensor_parallel_shards
    yield
    linear.bias = original_bias