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

# DeepSpeed Team

from abc import ABC, abstractmethod
from deepspeed.utils.types import ActivationFuncType, NormType
import torch
from deepspeed.accelerator import get_accelerator

transformer_param_names = (
        'attn_qkvw', \
        'attn_qkvb', \
        'attn_ow' , \
        'attn_ob', \
        'inter_w', \
        'inter_b', \
        'output_w', \
        'output_b', \
        'attn_nw', \
        'attn_nb', \
        'norm_w', \
        'norm_b')


class DSPolicy(ABC):
    _orig_layer_class = None

    def __init__(self):
        self.cuda_graph_supported = False

    @abstractmethod
    def attention(self):
        """
        Returns attention qkv and dense parameters
        weight: (3*hidden, hidden) and (hidden, hidden)
        bias: (3*hidden) and (hidden)
        """
        raise NotImplementedError


class TransformerPolicy(DSPolicy):
    # a static class variable containing the HuggingFace model configuration.
    # see e.g., transformers.models.opt.configuration_opt.OPTConfig
    hf_model_config = None

    def __init__(
            self,
            inference=True,
            linear_layer=True,
            scale_attention=True,
            megatron_v2=False,
            use_mup=False,
            # the type of activation function used in MLP
            mlp_act_func_type=ActivationFuncType.GELU,
            # applies layer norm before attention if `pre_attn_norm` is set to True
            pre_attn_norm=True,
            # this flag shows whether or not using prefix in loading the checkpoint
            use_load_prefix=False,
            # whether or not the qkv is stored in the split-format
            split_qkv=True,
            # Type of normalization to perform
            norm_type=NormType.LayerNorm):
        super().__init__()
        self.cuda_graph_supported = False
        self.inference = inference
        self.linear_layer = linear_layer
        self.scale_attention = scale_attention
        self.is_megatron_v2 = megatron_v2
        self.use_mup = use_mup
        self.mlp_act_func_type = mlp_act_func_type
        self.pre_attn_norm = pre_attn_norm
        self.use_load_prefix = use_load_prefix
        self.split_qkv = split_qkv
        self.norm_type = norm_type

    @abstractmethod
    def attention(self):
        """
        Returns attention qkv and dense parameters
        weight: (3*hidden, hidden) and (hidden, hidden)
        bias: (3*hidden) and (hidden)
        """
        raise NotImplementedError

    @abstractmethod
    def get_hidden_heads(self):
        """
        return hidden_size and number of heads
        """
        raise NotImplementedError

    @abstractmethod
    def mlp(self):
        """
        Returns mlp intermediate and output
        weight: (intermediate, hidden) and (hidden, intermediate)
        bias: (intermediate) and (hidden)
        """
        raise NotImplementedError

    @abstractmethod
    def layernorm(self):
        """
        Returns LayerNorms used in transformer layer
        Post-Attention and pre/post layer norm
        gamma and beta with shape: (hidden)
        """
        raise NotImplementedError


# TODO (lekurile): This function exists in base container as well, consolidate as some point
def transpose(data):
    with torch.no_grad():
        data = data.contiguous()
        data1 = data.transpose(-1, -2).reshape(-1)
        data.reshape(-1).copy_(data1)
        data1 = None
    return data.reshape(data.shape[-1], data.shape[-2])


# TODO (lekurile): This function exists in megatron feature container as well, consolidate as some point
def _transpose(x, heads=1, mp_replace=None):
    heads = heads // mp_replace.mp_size  # type: ignore
    outer_dim = -1
    attention_head_size = x.shape[outer_dim] // heads
    new_x_shape = x.size()[:outer_dim] + (heads, attention_head_size)
    x_1 = x.view(*new_x_shape)
    (q, k, v) = torch.split(x_1, (x_1.shape[-1] // 3), dim=-1)
    if len(q.shape) > 2:
        new_shape = (q.shape[0], ) + (-1, )
        return torch.cat((q.reshape(new_shape), k.reshape(new_shape), v.reshape(new_shape)),
                         dim=outer_dim).reshape(x.shape)
    else:
        return torch.cat((q.reshape(-1), k.reshape(-1), v.reshape(-1)), dim=-1).reshape(x.shape)


# This checks if the parameter exits in the checkpoint file and maybe copies it into the corresponding destination tensor.
# Note that not all parameters are saved in one checkpoint, that's why we always need to check if they exist!
def maybe_copy(module,
               sd,
               weight_quantizer,
               mp_replace,
               dst_name,
               src_name,
               qkv=False,
               megatron_v2=False,
               split_qkv=False,
               heads=1):
    if src_name in sd:
        dst = getattr(module, dst_name)
        tmp = sd[src_name]
        if len(dst.shape) == 1:
            if split_qkv:
                dst = mp_replace.strided_copy(dst, tmp, num_splits=3)
            else:
                dst = mp_replace.copy(dst, tmp)
            if qkv and megatron_v2:
                dst = torch.nn.parameter.Parameter(_transpose(dst, heads=heads, mp_replace=mp_replace).contiguous())
        else:
            if split_qkv:
                dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \
                                                (transpose(tmp).contiguous())), num_splits=3, int8=weight_quantizer.q_int8)
            else:
                if qkv and megatron_v2:
                    tmp = _transpose(transpose(tmp), heads=heads, mp_replace=mp_replace).contiguous()
                    if weight_quantizer.q_int8:
                        tmp = transpose(tmp)
                dst = mp_replace.copy(dst, weight_quantizer.quantize(tmp if weight_quantizer.q_int8 else \
                                                transpose(tmp)), int8=weight_quantizer.q_int8)
        setattr(module, dst_name, dst)


# Extending the maybe_copy function for when the q, k, and v are in separate parameters!
def maybe_copy_qkv(module, sd, weight_quantizer, mp_replace, dst_name, src_names, split_qkv=False):
    if src_names[0] in sd:
        q = sd[src_names[0]]
        k = sd[src_names[1]]
        v = sd[src_names[2]]
        qkv_data = torch.cat((q, k, v), dim=0)
        dst = getattr(module, dst_name)
        if len(dst.shape) == 1:
            if split_qkv:
                dst = mp_replace.strided_copy(dst, qkv_data.contiguous(), num_splits=3)
            else:
                dst = mp_replace.copy(dst, qkv_data)
        else:
            if split_qkv:
                dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \
                                                ((transpose(qkv_data)).contiguous())), num_splits=3, int8=weight_quantizer.q_int8)
            else:
                dst = mp_replace.copy(dst, weight_quantizer.quantize(qkv_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \
                                                transpose(qkv_data)), int8=weight_quantizer.q_int8)
        setattr(module, dst_name, dst)


# Extending the `maybe_copy` function for when mlp1 is in separate parameters for GeGLU
def maybe_copy_geglu(module, sd, weight_quantizer, mp_replace, dst_name, src_names):
    if src_names[0] in sd:
        reg_proj = sd[src_names[0]]
        gate_proj = sd[src_names[1]]

        mlp1_data = torch.cat((reg_proj, gate_proj), dim=0)
        dst = getattr(module, dst_name)

        dst = mp_replace.strided_copy(dst, weight_quantizer.quantize(mlp1_data.to(get_accelerator().device_name()) if weight_quantizer.q_int8 else \
                                            transpose(mlp1_data)), num_splits=2, int8=weight_quantizer.q_int8)
        setattr(module, dst_name, dst)


def pack_lora_weights(p):
    return [
        p.lora_right_weight, \
        p.lora_left_weight, \
        p.lora_scaling
    ]


def maybe_get_lora(p):
    if hasattr(p, 'lora_right_weight'):
        lora_param = pack_lora_weights(p)
    else:
        lora_param = []
    return lora_param