
import operator_benchmark as op_bench
import torch
import numpy as np


from torch.testing._internal.common_quantization import (
    lengths_to_offsets
)

torch.ops.load_library("//caffe2/torch/fb/sparsenn:sparsenn_operators")


embedding_bag_rowwise_offsets_short_configs = op_bench.cross_product_configs(
    num_embeddings=(80,),
    embedding_dim=(128, 256),
    num_offsets=range(2, 10),
    enable_per_sample_weights=(True, False),
    include_last_offset=(True, False),
    is_pruned_weights=(True, False,),
    use_32bit_indices=(True, False),
    use_32bit_offsets=(True, False),
    tags=('short',),
)


embedding_bag_rowwise_offsets_long_configs = op_bench.cross_product_configs(
    num_embeddings=(100, 120, 1000, 10_000, 20_000),
    embedding_dim=(16, 64, 128, 256),
    num_offsets=range(10, 20),
    enable_per_sample_weights=(True, False),
    include_last_offset=(True, False),
    is_pruned_weights=(True, False,),
    use_32bit_indices=(True, False),
    use_32bit_offsets=(True, False),
    tags=('long',)
)


full_configs = embedding_bag_rowwise_offsets_long_configs + embedding_bag_rowwise_offsets_short_configs

four_bit_rowwise_ops = op_bench.op_list(
    attrs=(
        ('qembeddingbag_4bit_rowwise_offsets', torch.ops.quantized.embedding_bag_4bit_rowwise_offsets),
    ),
    attr_names=('op_name', 'op_func'),
)

byte_rowwise_ops = op_bench.op_list(
    attrs=(
        ('qembeddingbag_byte_rowwise_offsets', torch.ops.quantized.embedding_bag_byte_rowwise_offsets),
    ),
    attr_names=('op_name', 'op_func'),
)


def get_pruned_weights_and_mapping(q_weights):
    indicator = torch.from_numpy(np.random.uniform(
        low=-1.0, high=1.0, size=[q_weights.shape[0]]).astype(np.float32))

    q_pruned_weights, compressed_indices_mapping = torch.ops.fb.embedding_bag_rowwise_prune(
        q_weights, indicator, 0.01, torch.int32)

    return q_pruned_weights, compressed_indices_mapping


class EmbedddingBag4BitRowwiseOffsetsTest(op_bench.TorchBenchmarkBase):
    def init(self,
             num_embeddings: int,
             embedding_dim: int,
             num_offsets: int,
             enable_per_sample_weights: bool,
             include_last_offset: bool,
             is_pruned_weights: bool,
             use_32bit_indices: bool,
             use_32bit_offsets: bool,
             op_func):
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.num_offsets = num_offsets
        self.enable_per_sample_weights = enable_per_sample_weights
        self.include_last_offset = include_last_offset
        self.max_segment_length = 20
        self.num_lengths = np.random.randint(1, num_offsets + 1)
        self.lengths = np.random.randint(0, self.max_segment_length + 1,
                                         size=self.num_lengths).astype(np.int32)
        self.num_indices = np.sum(self.lengths)
        self.is_pruned_weights = is_pruned_weights
        self.use_32bit_indices = use_32bit_indices
        self.use_32bit_offsets = use_32bit_offsets

        self.offsets = lengths_to_offsets(self.lengths)
        self.indices = torch.from_numpy(np.random.randint(
            low=0, high=num_embeddings, size=self.num_indices, dtype=np.int64))

        self.indices = self.indices.int() if self.use_32bit_indices else self.indices
        self.offsets = self.offsets.int() if self.use_32bit_offsets else self.offsets

        if self.include_last_offset:
            self.offsets = torch.cat(
                (self.offsets, torch.tensor([self.indices.size(0)], dtype=torch.long)), 0
            )

        self.weights = torch.from_numpy((np.random.random_sample((
            self.num_embeddings, self.embedding_dim)) + 1).astype(np.float32))
        self.indices = torch.from_numpy(np.random.randint(
            low=0, high=self.num_embeddings, size=self.num_indices, dtype=np.int64))
        self.prepack_func = torch.ops.quantized.embedding_bag_4bit_prepack

        self.prepacked_weights = self.prepack_func(self.weights)
        self.per_sample_weights = torch.from_numpy(np.random.uniform(
            low=0.01, high=0.5, size=[len(self.indices)]).astype(np.float32)) if \
            self.enable_per_sample_weights else None

        self.compressed_indices = None

        if self.is_pruned_weights:
            self.prepacked_weights, self.compressed_indices = get_pruned_weights_and_mapping(self.prepacked_weights)

        self.op_func = op_func

    def forward(self):
        return self.op_func(self.prepacked_weights, self.indices, self.offsets,
                            mode=0, per_sample_weights=self.per_sample_weights,
                            include_last_offset=self.include_last_offset,
                            pruned_weights=self.is_pruned_weights,
                            compressed_indices_mapping=self.compressed_indices)


class EmbedddingBagByteRowwiseOffsetsTest(op_bench.TorchBenchmarkBase):
    def init(self,
             num_embeddings: int,
             embedding_dim: int,
             num_offsets: int,
             enable_per_sample_weights: bool,
             include_last_offset: bool,
             is_pruned_weights: bool,
             use_32bit_indices: bool,
             use_32bit_offsets: bool,
             op_func):
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.num_offsets = num_offsets
        self.enable_per_sample_weights = enable_per_sample_weights
        self.include_last_offset = include_last_offset
        self.max_segment_length = 20
        self.num_lengths = np.random.randint(1, num_offsets + 1)
        self.lengths = np.random.randint(0, self.max_segment_length + 1,
                                         size=self.num_lengths).astype(np.int32)
        self.is_pruned_weights = is_pruned_weights
        self.use_32bit_indices = use_32bit_indices
        self.use_32bit_offsets = use_32bit_offsets

        self.num_indices = np.sum(self.lengths)
        self.offsets = lengths_to_offsets(self.lengths)
        self.indices = torch.from_numpy(np.random.randint(
            low=0, high=num_embeddings, size=self.num_indices, dtype=np.int64))

        self.indices = self.indices.int() if self.use_32bit_indices else self.indices
        self.offsets = self.offsets.int() if self.use_32bit_offsets else self.offsets

        if include_last_offset:
            self.offsets = torch.cat(
                (self.offsets, torch.tensor([self.indices.size(0)], dtype=torch.long)), 0
            )

        self.weights = torch.from_numpy((np.random.random_sample((
            self.num_embeddings, self.embedding_dim)) + 1).astype(np.float32))
        self.indices = torch.from_numpy(np.random.randint(
            low=0, high=self.num_embeddings, size=self.num_indices, dtype=np.int64))

        self.prepack_func = torch.ops.quantized.embedding_bag_byte_prepack

        self.prepacked_weights = self.prepack_func(self.weights)
        self.per_sample_weights = torch.from_numpy(np.random.uniform(
            low=0.01, high=0.5, size=[len(self.indices)]).astype(np.float32)) if \
            self.enable_per_sample_weights else None

        self.op_func = op_func

    def forward(self):
        return self.op_func(self.prepacked_weights, self.indices, self.offsets,
                            mode=0, per_sample_weights=self.per_sample_weights,
                            include_last_offset=self.include_last_offset,
                            pruned_weights=self.is_pruned_weights)


op_bench.generate_pt_tests_from_op_list(four_bit_rowwise_ops,
                                        full_configs,
                                        EmbedddingBag4BitRowwiseOffsetsTest)
op_bench.generate_pt_tests_from_op_list(byte_rowwise_ops,
                                        full_configs,
                                        EmbedddingBagByteRowwiseOffsetsTest)


if __name__ == "__main__":
    op_bench.benchmark_runner.main()
