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

# DeepSpeed Team

import torch
import os
import copy
import collections
import json
from abc import ABC, abstractmethod

from deepspeed.utils import logger
from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine

from .weight_quantizer import WeightQuantization

AUTO_MODULE_KEY = 'auto'


class SDLoaderFactory:

    @staticmethod
    def get_sd_loader_json(json_file, checkpoint_engine):
        if isinstance(json_file, str):
            with open(json_file) as f:
                data = json.load(f)
        else:
            assert isinstance(json_file, dict)
            data = json_file
        sd_type = data['type']
        ckpt_list = data['checkpoints']
        version = data['version']
        ckpt_type = data.get('parallelization', 'pp')
        mp_size = data.get('mp_size', 0)
        if sd_type.lower() in ['bloom', 'ds_model']:
            return data
        return SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine, sd_type, version)

    @staticmethod
    def get_sd_loader(ckpt_list, checkpoint_engine, sd_type='Megatron', version=None):
        if sd_type == 'Megatron':
            return MegatronSDLoader(ckpt_list, version, checkpoint_engine)
        else:
            assert False, '{} checkpoint type is not supported'.format(sd_type)


class SDLoaderBase(ABC):

    def __init__(self, ckpt_list, version, checkpoint_engine):
        self.module_key = None
        self.ckpt_list = ckpt_list
        self.version = version
        self.checkpoint_engine = TorchCheckpointEngine() if checkpoint_engine is None else checkpoint_engine
        self.check_ckpt_list()

    def load(self,
             mp_world_size,
             mp_rank,
             module_key=AUTO_MODULE_KEY,
             is_pipe_parallel=False,
             quantize=False,
             quantize_bits=8,
             quantize_groups=64,
             mlp_extra_grouping=True):
        self.module_key = module_key
        num_ckpt = len(self.ckpt_list)
        idx = mp_rank * num_ckpt // mp_world_size
        """ We have multiple cases to handle here for both training and inference:
            1. PipeModule loading mp_rank_*.pt files, is_pipe_parallel=True, module_key is not None
                a. if no mp_size/pp_size resizing occurs, for both training & inference, loading
                   the mp_rank related checkpoint directly.
                b. if has mp_size/pp_size resizing, only Megatron model inference is supported,
                   in this case each mp_rank_*.pt have same content, we will load the first checkpoint
                   file (idx=0), to avoid idx exceeding file list boundary.

            2. PipeModule loading layer_*.pt files, is_pipe_parallel=True, module_key is None
                a. if no mp_size resizing occurs, for both training & inference, loading
                   the mp_rank related checkpoint directly.
                b. if has mp_size resizing, only Megatron model inference is supported,
                   checkpoint file(s) will be merged/split according to mp_rank, mp_world_size and
                   checkpoint file list.

            3. Non-PipeModule loading mp_rank_*.pt files, is_pipe_parallel=False
                Same with case (2).
        """
        if is_pipe_parallel and module_key is not None and mp_world_size != num_ckpt:
            mp_world_size = num_ckpt
            idx = 0

        load_path = self.ckpt_list[idx]

        merge_count = 1
        if num_ckpt == mp_world_size:
            assert os.path.exists(load_path)
            #logger.info(f'rank: {mp_rank} loading checkpoint: {load_path}')
            sd = self.checkpoint_engine.load(load_path, map_location=lambda storage, \
                loc: storage)

            if quantize:
                quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size)
                sd_module, all_scales = quantizer.sd_quantize_megatron(self.get_module(sd), quantize_bits,
                                                                       quantize_groups)
                self.set_module(sd, sd_module)
            else:
                all_scales = None
        elif num_ckpt > mp_world_size:
            sd, all_scales, merge_count = self.merge_state_dict(mp_world_size, mp_rank, quantize, \
                quantize_bits, quantize_groups, mlp_extra_grouping)
        else:
            sd, all_scales = self.split_state_dict(mp_world_size, mp_rank, quantize, quantize_bits, \
                quantize_groups, mlp_extra_grouping)
        return load_path, sd, (all_scales, merge_count)

    def get_merge_state_dicts(self, mp_world_size, mp_rank):
        num_ckpt = len(self.ckpt_list)
        assert num_ckpt % mp_world_size == 0, 'Invalid checkpoints and world size for sd merge'

        num_to_merge = num_ckpt // mp_world_size
        ckpt_list = [self.ckpt_list[i] for i in range(num_to_merge * mp_rank, num_to_merge * (mp_rank + 1))]

        logger.info(f"mp_rank: {mp_rank}, ckpt_list: {ckpt_list}")
        sd_list = [self.checkpoint_engine.load(ckpt, map_location=lambda storage, loc: storage) for ckpt in ckpt_list]
        return sd_list

    def get_split_state_dict(self, mp_world_size, mp_rank):
        num_ckpt = len(self.ckpt_list)
        assert mp_world_size % num_ckpt == 0, 'Invalid checkpoints and world size for sd split'

        num_to_split = mp_world_size // num_ckpt
        ckpt_index = mp_rank // num_to_split
        ckpt_offset = mp_rank % num_to_split

        logger.info(f"mp_rank: {mp_rank}, ckpt_list: {self.ckpt_list[ckpt_index]}, offset: {ckpt_offset}")

        sd = self.checkpoint_engine.load(self.ckpt_list[ckpt_index], map_location=lambda storage, loc: storage)

        return sd, num_to_split, ckpt_offset

    def _choose_module_key(self, sd):
        assert not ('module' in sd
                    and 'model' in sd), "checkpoint has both 'model' and 'module' keys, not sure how to proceed"
        assert 'module' in sd or 'model' in sd, "checkpoint contains neither 'model' or 'module' keys, not sure how to proceed"
        if 'module' in sd:
            return 'module'
        elif 'model' in sd:
            return 'model'

    def get_module(self, sd):
        if self.module_key is None:
            return sd
        elif self.module_key == AUTO_MODULE_KEY:
            return sd[self._choose_module_key(sd)]
        else:
            return sd[self.module_key]

    def set_module(self, sd, module):
        if self.module_key is None:
            sd = module
        elif self.module_key == AUTO_MODULE_KEY:
            sd[self._choose_module_key(sd)] = module
        else:
            sd[self.module_key] = module
        return sd

    def check_ckpt_list(self):
        #logger.info(f'checkpoint file list: {self.ckpt_list}')
        assert len(self.ckpt_list) > 0

        sd = self.checkpoint_engine.load(self.ckpt_list[0], map_location=lambda storage, loc: storage)

        # check checkpoint count is same with saved mp_world_size
        if 'mp_world_size' in sd.keys():
            assert len(self.ckpt_list) == sd[
                'mp_world_size'], f"checkpoint count {len(self.ckpt_list)} is different from saved mp_world_size {sd['mp_world_size']}"

    @abstractmethod
    def merge_state_dict(self, mp_world_size, mp_rank, quantize, quantize_bits, groups, mlp_extra_grouping):
        pass

    @abstractmethod
    def split_state_dict(self, mp_world_size, mp_rank, quantize, quantize_bits, groups, mlp_extra_grouping):
        pass

    @abstractmethod
    def sanity_check(self, ckpt_file_name):
        pass


class MegatronSDLoader(SDLoaderBase):

    def __init__(self, ckpt_list, version, checkpoint_engine):
        super().__init__(ckpt_list, version, checkpoint_engine)
        """
        ## Q/K/V data need special processing
        key: transformer.layers.0.attention.query_key_value.weight, shape: torch.Size([3192, 4256])
        key: transformer.layers.0.attention.query_key_value.bias, shape: torch.Size([3192])

        ## merge or split on axis=0
        key: word_embeddings.weight, shape: torch.Size([12672, 4256])
        key: transformer.layers.0.mlp.dense_h_to_4h.bias, shape: torch.Size([4256])
        key: transformer.layers.0.mlp.dense_h_to_4h.weight, shape: torch.Size([4256, 4256])

        ## merge or split on axis=1
        key: transformer.layers.0.attention.dense.weight, shape: torch.Size([4256, 1064])
        key: transformer.layers.0.mlp.dense_4h_to_h.weight, shape: torch.Size([4256, 4256])

        ## no change required
        key: transformer.layers.0.mlp.dense_4h_to_h.bias, shape: torch.Size([4256])
        key: transformer.final_layernorm.weight, shape: torch.Size([4256])
        key: transformer.final_layernorm.bias, shape: torch.Size([4256])
        key: transformer.layers.0.attention.dense.bias, shape: torch.Size([4256])
        key: transformer.layers.0.post_attention_layernorm.weight, shape: torch.Size([4256])
        key: transformer.layers.0.post_attention_layernorm.bias, shape: torch.Size([4256])
        key: transformer.layers.0.input_layernorm.weight, shape: torch.Size([4256])
        key: transformer.layers.0.input_layernorm.bias, shape: torch.Size([4256])
        key: position_embeddings.weight, shape: torch.Size([1024, 4256])
        """

    def merge_query_key_value(self, param_list, ckpt_ver):
        """
        Up to now we found 3 Q/K/V parameter formats in different Megatron checkpoint versions:

        1. version 0, there is no version information saved in checkpoint.
            format: [(3 * np * hn), h]
        2. version 1.0
            format: [(np * hn * 3), h]
        3. version 2.0
            format: [(np * 3 * hn), h]

        h: hidden size
        n: number of attention heads
        p: number of model parallel partitions
        np: n/p
        hn: h/n
        """

        new_qkv = None
        if ckpt_ver == 0:
            # [(3 * np * hn), h]
            assert param_list[0].shape[0] % 3 == 0
            size_qkv = param_list[0].shape[0] // 3
            split_tensors = [torch.split(param, size_qkv, dim=0) for param in param_list]

            tensors = []
            for i in range(3):
                tensor_tuple = [t[i] for t in split_tensors]
                tensors.append(torch.cat(tensor_tuple, axis=0))
            new_qkv = torch.cat(tensors, axis=0)
        elif ckpt_ver == 1.0 or ckpt_ver == 2.0:
            # [(np * hn * 3), h] or [(np * 3 * hn), h]
            new_qkv = torch.cat(param_list, axis=0)
        else:
            assert False, f'checkpoint version: {ckpt_ver} is not supported'

        return new_qkv

    def split_query_key_value(self, param, num_to_split, offset, ckpt_ver):
        """
        Up to now we found 3 Q/K/V parameter formats in different Megatron checkpoint versions:

        1. version 0, there is no version information saved in checkpoint.
            format: [(3 * np * hn), h]
        2. version 1.0
            format: [(np * hn * 3), h]
        3. version 2.0
            format: [(np * 3 * hn), h]

        h: hidden size
        n: number of attention heads
        p: number of model parallel partitions
        np: n/p
        hn: h/n
        """

        new_qkv = None
        if ckpt_ver == 0:
            # [(3 * np * hn), h]
            assert param.shape[0] % 3 == 0
            size_qkv = param.shape[0] // 3
            split_tensors = torch.split(param, size_qkv, dim=0)

            assert split_tensors[0].shape[0] % num_to_split == 0
            split_size = split_tensors[0].shape[0] // num_to_split

            tensors = []
            for i in range(3):
                tensors.append(torch.split(split_tensors[i], split_size, dim=0)[offset])
            new_qkv = torch.cat(tensors, axis=0)
        elif ckpt_ver == 1.0 or ckpt_ver == 2.0:
            # [(np * hn * 3), h] or [(np * 3 * hn), h]
            assert param.shape[0] % num_to_split == 0
            size_qkv = param.shape[0] // num_to_split
            split_tensors = torch.split(param, size_qkv, dim=0)
            new_qkv = split_tensors[offset]
        else:
            assert False, f'checkpoint version: {ckpt_ver} is not supported'

        return new_qkv

    def merge_state_dict(self,
                         mp_world_size,
                         mp_rank,
                         quantize=False,
                         quantize_bits=8,
                         groups=64,
                         mlp_extra_grouping=True):
        self.sanity_check(self.ckpt_list[0])

        sd_list = self.get_merge_state_dicts(mp_world_size, mp_rank)
        ds_sd = copy.deepcopy(sd_list[0])
        new_client_sd = collections.OrderedDict()

        client_sd_list = [self.get_module(sd) for sd in sd_list]
        keys = client_sd_list[0].keys()

        ckpt_ver = self.get_checkpoint_version(ds_sd)
        logger.info(f"checkpoint version: {ckpt_ver}")
        if quantize:
            quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size)

        for key in keys:
            value_list = [sd[key] for sd in client_sd_list]

            if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key:
                if quantize:
                    value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key, merge_dim=1)
                new_client_sd[key] = torch.cat(value_list, axis=1)
            elif "attention.query_key_value" in key:
                if quantize and "attention.query_key_value.weight" in key:
                    value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key)
                    new_client_sd[key] = torch.cat(value_list, axis=0)
                else:
                    if quantize:
                        new_client_sd[key] = torch.cat(value_list, axis=0)
                    else:
                        new_client_sd[key] = self.merge_query_key_value(value_list, ckpt_ver)
            elif "mlp.dense_h_to_4h.weight" in key or "word_embeddings.weight" in key or "mlp.dense_h_to_4h.bias" in key:
                if quantize and "mlp.dense_h_to_4h.weight" in key:
                    value_list = quantizer.Quantize(value_list, quantize_bits, groups, key=key)
                new_client_sd[key] = torch.cat(value_list, axis=0)
            else:
                new_client_sd[key] = value_list[0]
        if quantize:
            all_scales = quantizer.merge_scales()
        ds_sd = self.set_module(ds_sd, new_client_sd)

        return ds_sd, (all_scales if quantize else None), len(client_sd_list)

    def split_state_dict(self,
                         mp_world_size,
                         mp_rank,
                         quantize=False,
                         quantize_bits=8,
                         groups=64,
                         mlp_extra_grouping=True):
        #self.sanity_check(self.ckpt_list[0])

        sd, num_to_split, ckpt_offset = self.get_split_state_dict(mp_world_size, mp_rank)
        ds_sd = copy.deepcopy(sd)
        new_client_sd = collections.OrderedDict()

        client_sd = self.get_module(sd)

        ckpt_ver = self.get_checkpoint_version(ds_sd)
        logger.info(f"checkpoint version: {ckpt_ver}")

        if quantize:
            quantizer = WeightQuantization(mlp_extra_grouping=mlp_extra_grouping, mp_size=mp_world_size)

        for key in client_sd.keys():
            value = client_sd[key]

            if "attention.dense.weight" in key or "mlp.dense_4h_to_h.weight" in key:
                assert value.shape[1] % num_to_split == 0
                split_size = value.shape[1] // num_to_split
                if quantize:
                    q_vals = quantizer.Quantize([value], quantize_bits, groups, key)
                    value = q_vals[0]
                new_client_sd[key] = torch.split(value, split_size, dim=1)[ckpt_offset]
            elif "attention.query_key_value" in key:
                if quantize and "attention.query_key_value.weight" in key:
                    q_vals = quantizer.Quantize([value], quantize_bits, groups, key)
                    value = q_vals[0]
                new_client_sd[key] = self.split_query_key_value(value, num_to_split, ckpt_offset, ckpt_ver)
            elif "mlp.dense_h_to_4h.weight" in key or "word_embeddings.weight" in key or "mlp.dense_h_to_4h.bias" in key or "final_linear.weight" in key:
                assert value.shape[0] % num_to_split == 0
                split_size = value.shape[0] // num_to_split
                if quantize and "mlp.dense_h_to_4h.weight" in key:
                    q_vals = quantizer.Quantize([value], quantize_bits, groups, key)
                    value = q_vals[0]
                new_client_sd[key] = torch.split(value, split_size, dim=0)[ckpt_offset]
            else:
                new_client_sd[key] = value

        if quantize:
            all_scales = quantizer.merge_scales_split(num_to_split)

        ds_sd = self.set_module(ds_sd, new_client_sd)

        return ds_sd, (all_scales if quantize else None)

    def sanity_check(self, ckpt_file_name):
        keys_to_check = [
            "attention.dense.weight", "mlp.dense_4h_to_h.weight", "attention.query_key_value",
            "mlp.dense_h_to_4h.weight", "mlp.dense_h_to_4h.bias"
        ]

        sd = self.checkpoint_engine.load(ckpt_file_name, map_location=lambda storage, loc: storage)

        # partial_key is a sub-string of one key in the sd
        def check_key_exist(partial_key, sd):
            keys = sd.keys()
            found = False
            for k in keys:
                if partial_key in k:
                    found = True
                    break
            return found

        for key in keys_to_check:
            assert check_key_exist(key,
                                   self.get_module(sd)), f'key: {key} is not found in the checkpoint {ckpt_file_name}'

    def get_checkpoint_version(self, state_dict):
        # Use 0 if version info doesn't exist
        return self.version if self.version is not None else state_dict.get('checkpoint_version', 0)