from typing import List, Optional, Tuple

import torch
import torch.nn.functional as F

from colossalai.legacy.nn._ops._utils import dual_all_to_all
from colossalai.legacy.tensor import ColoTensorSpec, ComputePattern, ProcessGroup, ShardSpec
from colossalai.tensor import ColoTensor

from .cache_mgr import EvictionStrategy
from .cached_embedding import CachedEmbeddingBag


def get_partition(embedding_dim, rank, world_size) -> Tuple[int, int, bool]:
    if world_size == 1:
        return 0, embedding_dim, True

    assert embedding_dim >= world_size, (
        f"Embedding dimension {embedding_dim} must be larger than the world size " f"{world_size} of the process group"
    )
    chunk_size = embedding_dim // world_size
    threshold = embedding_dim % world_size
    # if embedding dim is divisible by world size
    if threshold == 0:
        return rank * chunk_size, (rank + 1) * chunk_size, True

    # align with the split strategy of torch.tensor_split
    size_list = [chunk_size + 1 if i < threshold else chunk_size for i in range(world_size)]
    offset = sum(size_list[:rank])
    return offset, offset + size_list[rank], False


class ParallelCachedEmbeddingBag(CachedEmbeddingBag):
    def __init__(
        self,
        num_embeddings,
        embedding_dim,
        padding_idx=None,
        max_norm=None,
        norm_type=2.0,
        scale_grad_by_freq=False,
        sparse=False,
        _weight=None,
        mode="mean",
        include_last_offset=False,
        dtype=None,
        device=None,
        cache_ratio=0.01,
        ids_freq_mapping=None,
        warmup_ratio=0.7,
        buffer_size=50_000,
        pin_weight=False,
        evict_strategy: EvictionStrategy = EvictionStrategy.DATASET,
    ):
        self.rank = torch.distributed.get_rank()
        self.world_size = torch.distributed.get_world_size()

        self.partition_start_index, self.partition_end_index, divisible = get_partition(
            embedding_dim, self.rank, self.world_size
        )
        self.embedding_dim_per_partition = self.partition_end_index - self.partition_start_index

        super(ParallelCachedEmbeddingBag, self).__init__(
            num_embeddings,
            embedding_dim,
            padding_idx,
            max_norm,
            norm_type,
            scale_grad_by_freq,
            sparse,
            _weight,
            mode,
            include_last_offset,
            dtype,
            device,
            cache_ratio,
            ids_freq_mapping,
            warmup_ratio,
            buffer_size,
            pin_weight,
            evict_strategy,
        )
        self.cache_op = True

    def _weight_alloc(self, dtype, device):
        weight = torch.empty(self.num_embeddings, self.embedding_dim_per_partition, device=device, dtype=dtype)
        with torch.no_grad():
            weight.data.uniform_(-1 / self.num_embeddings, 1 / self.num_embeddings)
            if self.padding_idx is not None:
                weight[self.padding_idx].fill_(0)
        colo_tensor_spec = ColoTensorSpec(
            pg=ProcessGroup(tp_degree=self.world_size),
            dist_attr=ShardSpec(dims=[-1], num_partitions=[self.world_size]),
            compute_attr=ComputePattern.TP1D,
        )
        return ColoTensor.from_torch_tensor(weight, spec=colo_tensor_spec)

    def forward(
        self,
        indices,
        offsets=None,
        per_sample_weights=None,
        shape_hook=None,
        scatter_dim=0,
        gather_dim=-1,
    ):
        if self.cache_op:
            with torch.no_grad():
                indices = self.cache_weight_mgr.prepare_ids(indices)
        output_shard = F.embedding_bag(
            indices.cuda(),
            self.cache_weight_mgr.cuda_cached_weight,
            offsets,
            self.max_norm,
            self.norm_type,
            self.scale_grad_by_freq,
            self.mode,
            self.sparse,
            per_sample_weights,
            self.include_last_offset,
            self.padding_idx,
        )
        if shape_hook is not None:
            output_shard = shape_hook(output_shard)
        output_full = dual_all_to_all(
            output_shard, self.weight.get_process_group(), scatter_dim=scatter_dim, gather_dim=gather_dim
        )
        return output_full

    def set_cache_op(self, cache_op: bool = True):
        self.cache_op = cache_op

    @classmethod
    def from_pretrained(
        cls,
        embedding: torch.Tensor,
        freeze: bool = True,
        padding_idx: Optional[int] = None,
        max_norm: Optional[float] = None,
        norm_type: float = 2.0,
        scale_grad_by_freq: bool = False,
        sparse: bool = False,
        mode: str = "mean",
        include_last_offset: bool = False,
        cuda_row_num: int = 100_000,
        ids_freq_mapping: Optional[List[int]] = None,
        warmup_ratio: float = 0.7,
        buffer_size: int = 0,
    ) -> "ParallelCachedEmbeddingBag":
        rows, cols = embedding.shape
        embedding_bag = cls(
            rows,
            cols,
            padding_idx,
            max_norm,
            norm_type,
            scale_grad_by_freq,
            sparse,
            embedding,
            mode,
            include_last_offset,
            cuda_row_num=cuda_row_num,
            ids_freq_mapping=ids_freq_mapping,
            warmup_ratio=warmup_ratio,
            buffer_size=buffer_size,
        )
        embedding_bag.cache_weight_mgr.cuda_cached_weight.requires_grad_ = not freeze
        return embedding_bag

    def print_comm_stats_(self):
        self.cache_weight_mgr.print_comm_stats()

    def element_size(self):
        return self.weight.element_size()
