# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
import pytest
import random
import numpy as np
from unit.megatron_model import get_gpt2_model
from deepspeed.compression.compress import init_compression
from unit.modeling import BertConfig
from unit.modelingpreln import BertEncoder as BertEncoderPreln
from deepspeed.compression.basic_layer import LinearLayer_Compress, ColumnParallelLinear_Compress, RowParallelLinear_Compress
from deepspeed.compression.helper import convert_conv1d_to_linear
from deepspeed.accelerator import get_accelerator
from deepspeed.utils.torch import required_torch_version
from unit.common import DistributedTest

pytestmark = pytest.mark.skipif(not required_torch_version(min_version=1.5),
                                reason='Megatron-LM package requires Pytorch version 1.5 or above')


def reset_random(seed=1234):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    get_accelerator().manual_seed_all(seed)


def create_bert_model():
    hidden_size = 384
    num_layers = 2
    heads = 12
    dropout_ratio = 0.1
    bert_config = BertConfig(vocab_size_or_config_json_file=119547,
                             hidden_size=hidden_size,
                             num_hidden_layers=num_layers,
                             num_attention_heads=heads,
                             intermediate_size=hidden_size * 4,
                             hidden_act="gelu",
                             hidden_dropout_prob=dropout_ratio,
                             attention_probs_dropout_prob=dropout_ratio,
                             max_position_embeddings=512,
                             type_vocab_size=2,
                             initializer_range=0.2)

    weights = []
    biases = []

    for i in range(4):
        weights.append(torch.nn.Parameter(torch.Tensor(hidden_size, hidden_size)))

    weights.append(torch.nn.Parameter(torch.Tensor(hidden_size)))
    weights.append(torch.nn.Parameter(torch.Tensor(hidden_size * 4, hidden_size)))
    weights.append(torch.nn.Parameter(torch.Tensor(hidden_size, hidden_size * 4)))
    weights.append(torch.nn.Parameter(torch.Tensor(hidden_size)))

    biases.append(torch.nn.Parameter(torch.Tensor(hidden_size)))
    for i in range(4):
        biases.append(torch.nn.Parameter(torch.Tensor(hidden_size)))
    biases.append(torch.nn.Parameter(torch.Tensor(hidden_size * 4)))
    biases.append(torch.nn.Parameter(torch.Tensor(hidden_size)))
    biases.append(torch.nn.Parameter(torch.Tensor(hidden_size)))

    return BertEncoderPreln(bert_config, weights, biases)


class Conv1D(torch.nn.Module):
    """
    1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2).
    Basically works like a linear layer but the weights are transposed.
    Args:
        nf (`int`): The number of output features.
        nx (`int`): The number of input features.
    """

    def __init__(self, nf, nx):
        super().__init__()
        self.nf = nf
        w = torch.empty(nx, nf)
        self.weight = torch.nn.Parameter(w)
        self.bias = torch.nn.Parameter(torch.zeros(nf))

    def forward(self, x):
        size_out = x.size()[:-1] + (self.nf, )
        x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
        x = x.view(size_out)
        return x


def create_conv1d_model():
    nf = 128
    nx = 128

    return torch.nn.ModuleList([Conv1D(nf, nx) for i in range(4)])


class TestCompression(DistributedTest):

    def setup_method(self, method):
        reset_random()

    def get_ds_config(self):
        ds_config_dict = {
            "train_micro_batch_size_per_gpu": 1,
            "optimizer": {
                "type": "Lamb",
                "params": {
                    "lr": 0.00015
                }
            },
            "fp16": {
                "enabled": True
            },
            "compression_training": {
                "weight_quantization": {
                    "shared_parameters": {
                        "enabled": True,
                        "quantizer_kernel": False,
                        "schedule_offset": 50,
                        "quantize_groups": 1,
                        "quantize_verbose": False,
                        "quantization_type": "asymmetric",
                        "rounding": "nearest",
                        "fp16_mixed_quantize": {
                            "enabled": False,
                            "quantize_change_ratio": 0.001
                        }
                    },
                    "different_groups": {
                        "wq1": {
                            "params": {
                                "start_bits": 12,
                                "target_bits": 8,
                                "quantization_period": 50
                            },
                            "modules": ["attention.self", "intermediate"]
                        },
                        "wq2": {
                            "params": {
                                "start_bits": 12,
                                "target_bits": 4,
                                "quantization_period": 50
                            },
                            "modules": ["attention.output"]
                        }
                    }
                },
                "activation_quantization": {
                    "shared_parameters": {
                        "enabled": True,
                        "quantization_type": "asymmetric",
                        "range_calibration": "dynamic",
                        "schedule_offset": 50
                    },
                    "different_groups": {
                        "aq1": {
                            "params": {
                                "bits": 8
                            },
                            "modules": ["attention.output"]
                        }
                    }
                },
                "sparse_pruning": {
                    "shared_parameters": {
                        "enabled": True,
                        "schedule_offset": 30,
                        "method": "l1"
                    },
                    "different_groups": {
                        "sp1": {
                            "params": {
                                "dense_ratio": 0.5
                            },
                            "modules": ["attention.self"]
                        }
                    }
                },
                "row_pruning": {
                    "shared_parameters": {
                        "enabled": True,
                        "schedule_offset": 20,
                        "method": "topk"
                    },
                    "different_groups": {
                        "rp1": {
                            "params": {
                                "dense_ratio": 0.5
                            },
                            "modules": ["intermediate.dense"],
                            "related_modules": [["layer.\\w+.output.dense"]]
                        }
                    }
                },
                "head_pruning": {
                    "shared_parameters": {
                        "enabled": True,
                        "schedule_offset": 10,
                        "method": "topk",
                        "num_heads": 12
                    },
                    "different_groups": {
                        "rp1": {
                            "params": {
                                "dense_ratio": 0.5
                            },
                            "modules": ["attention.output.dense"],
                            "related_modules": [["self.query", "self.key", "self.value"]]
                        }
                    }
                }
            }
        }

        return ds_config_dict

    def test_linear_layer_compress(self, tmpdir):
        model = create_bert_model()
        compressed_model = init_compression(model, self.get_ds_config())

        assert isinstance(compressed_model.layer[0].attention.self.query, LinearLayer_Compress)
        assert isinstance(compressed_model.layer[0].attention.self.key, LinearLayer_Compress)
        assert isinstance(compressed_model.layer[0].attention.self.value, LinearLayer_Compress)

    @pytest.mark.skip(reason="megatron-lm is currently broken so this test cannot be run.")
    def test_mpu_compress(self, tmpdir):
        if not required_torch_version(max_version=1.13):
            pytest.skip("megatron not compatible with torch >1.13")
        from megatron import mpu
        args_defaults = {
            'num_layers': 2,
            'hidden_size': 128,
            'num_attention_heads': 8,
            'max_position_embeddings': 128,
        }

        model = get_gpt2_model(args_defaults)
        compressed_model = init_compression(model, self.get_ds_config(), mpu=mpu)

        assert isinstance(compressed_model.module.language_model.transformer.layers[0].attention.query_key_value,
                          ColumnParallelLinear_Compress)
        assert isinstance(compressed_model.module.language_model.transformer.layers[0].attention.dense,
                          RowParallelLinear_Compress)
        assert isinstance(compressed_model.module.language_model.transformer.layers[0].mlp.dense_h_to_4h,
                          ColumnParallelLinear_Compress)
        assert isinstance(compressed_model.module.language_model.transformer.layers[0].mlp.dense_4h_to_h,
                          RowParallelLinear_Compress)

    def test_conv1d_convertion(self, tmpdir):
        model = create_conv1d_model()
        compressed_model = convert_conv1d_to_linear(model, Conv1D)

        assert isinstance(compressed_model[0], torch.nn.Linear)
        assert isinstance(compressed_model[1], torch.nn.Linear)
        assert isinstance(compressed_model[2], torch.nn.Linear)
        assert isinstance(compressed_model[3], torch.nn.Linear)