# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import math
import re
import shutil
from collections import OrderedDict

import paddle


class ParallelConfig:
    def __init__(self, mp: int, pp: int, vpp: int = 1, sharding: int = 1):
        self.mp = mp
        self.pp = pp
        self.vpp = vpp
        self.sharding = sharding

    def pipe_parallel_group(self, i: int, j: int):
        ans = []
        for k in range(self.pp):
            ans.append((i, j, k))
        return ans


class LayerReNamingHelper:
    def __init__(self, template: str):
        self._template = template
        self._i = -1
        self._last_old_layer_name = None

    def get_new_layer_name(self, old_layer_name: str):
        old_layer_name = old_layer_name.split(".")[0]
        if (
            self._last_old_layer_name is None
            or old_layer_name != self._last_old_layer_name
        ):
            self._i = self._i + 1
            self._last_old_layer_name = old_layer_name
        return self._template.format(self._i)


class LayerReNamingManager:
    def __init__(self):
        self._renaming_helpers = OrderedDict()
        self._renaming_helpers["linear"] = LayerReNamingHelper("linear_{}")
        self._renaming_helpers["layer_norm"] = LayerReNamingHelper(
            "layer_norm_{}"
        )
        self._renaming_helpers["embedding"] = LayerReNamingHelper(
            "embedding_{}"
        )

    def get_new_layer_name(self, old_name: str):
        layer_name = ""
        for k, v in self._renaming_helpers.items():
            if old_name.startswith(k):
                layer_name = v.get_new_layer_name(old_name)
                break
        return layer_name

    def get_new_param_name(self, old_name: str):
        names = old_name.split(".")
        layer_name = self.get_new_layer_name(names[0])
        assert layer_name, f"can not rename layer {names[0]}"
        names[0] = layer_name
        return ".".join(names)


class PipeLineModelAdaptor:
    def __init__(
        self,
        src_parallel_config: ParallelConfig,
        dst_parallel_config: ParallelConfig,
        transformer_layer_num: int,
        segment_method: str = "layer",
    ):
        self._src_parallel_config = src_parallel_config
        self._dst_parallel_config = dst_parallel_config
        self._transformer_layer_num = transformer_layer_num
        self._segment_method = segment_method

    def apply(self, src_model_path: str, dst_model_path: str):
        for i in range(self._src_parallel_config.mp):
            for j in range(self._src_parallel_config.sharding):
                # TODO(liuzhenhai): use multiple process
                layers = []

                # 1、extract layers in the same pp group
                group = self._src_parallel_config.pipe_parallel_group(i, j)
                src_dirs = [
                    "{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(
                        src_model_path, *e
                    )
                    for e in group
                ]
                # first rank extract shared layer
                with_shared = True
                for dir in src_dirs:
                    print(f"extract layer params in dir {dir}")
                    layers.extend(self.extract_layers(dir, with_shared))
                    with_shared = False
                # 2、sort and unique layers
                layers = self.sort_layers(layers)

                # 3、resplit layers among pp group according new pp config
                layer_segments = self.segment_layers(
                    layers, self._dst_parallel_config, self._segment_method
                )
                dst_group = self._dst_parallel_config.pipe_parallel_group(i, j)
                dst_dirs = [
                    "{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(
                        dst_model_path, *e
                    )
                    for e in dst_group
                ]

                # 4、merge layers belonging to the same node
                for layer_segment, dir_ in zip(layer_segments, dst_dirs):
                    print(f"merge {len(layer_segment)} layers to {dir_}")
                    self.merge_layers(layer_segment, dir_)

                # 5、copy meta_state.pdopt
                for src_dir, dst_dir in zip(src_dirs, dst_dirs):
                    shutil.copyfile(
                        f"{src_dir}/meta_state.pdopt",
                        f"{dst_dir}/meta_state.pdopt",
                    )

    def peek_model(self, model_dir: str):
        for i in range(self._src_parallel_config.mp):
            for j in range(self._src_parallel_config.sharding):
                group = self._src_parallel_config.pipe_parallel_group(i, j)
                dirs = [
                    "{}/mp_{:0>2d}_sharding_{:0>2d}_pp_{:0>2d}".format(
                        model_dir, *e
                    )
                    for e in group
                ]
                for dir in dirs:
                    print(f"peek partial model in {dir}:")
                    self.peek_partial_model(dir)

    def peek_partial_model(self, sub_dir: str):
        state_dict = paddle.load(f"{sub_dir}/model.pdparams")
        for k, v in state_dict.items():
            print(f"\t{k} -> {v.name}")

    def extract_layers(self, dir: str, with_shared: bool):
        opt = paddle.load(dir + "/model_state.pdopt")
        params = paddle.load(dir + "/model.pdparams")
        shared_layer_parsed = False
        # tname -> (layer, param_name)
        tname_to_layer_and_pname = {}
        for k, v in params.items():
            layer = self._extract_layer_name(k)
            assert layer
            # special treatment for embedding layer, skip duplicated shared layer
            # shared layer may exist or not, if it exist it share weight with _layers.0
            # _layers.shared_layers.embed.word_embeddings.weight -> embedding_0.w_0
            # _layers.shared_layers.embed.position_embeddings.weight -> embedding_1.w_0
            # _layers.0.word_embeddings.weight -> embedding_0.w_0
            # _layers.0.position_embeddings.weight -> embedding_1.w_0
            shared_layer_parsed = shared_layer_parsed or (
                "_layers.shared_layers" in layer
            )
            if (
                "_layers.shared_layers" not in layer
                and ("word_embeddings" in k or "position_embeddings" in k)
                and shared_layer_parsed
            ):
                continue
            tname_to_layer_and_pname[v.name] = (layer, k)

            # get opt-> param mapping
        tensor_names = list(tname_to_layer_and_pname.keys())
        opt_names = [
            e for e in opt.keys() if e not in ["master_weights", "LR_Scheduler"]
        ]
        opt_to_t = self._opt_name_to_tname(tensor_names, opt_names)
        # gather tensors belonging to one layer together
        layers = OrderedDict()
        for k, v in params.items():
            layer, p = tname_to_layer_and_pname[v.name]
            if layer not in layers:
                layers[layer] = {}
                layers[layer]["opt"] = OrderedDict()
                layers[layer]["params"] = OrderedDict()
                layers[layer]["master_weights"] = OrderedDict()
            layers[layer]["params"][p] = v

        for k, v in opt.items():
            if k in ["master_weights", "LR_Scheduler"]:
                continue
            layer, _ = tname_to_layer_and_pname[opt_to_t[v.name]]
            layers[layer]["opt"][k] = v

        if "master_weights" in opt:
            for k, v in opt["master_weights"].items():
                layer, _ = tname_to_layer_and_pname[k]
                layers[layer]["master_weights"][k] = v

        if "LR_Scheduler" in opt:
            for layer in layers:
                layers[layer]["LR_Scheduler"] = opt["LR_Scheduler"]

        ans = []

        for layer_name, layer in layers.items():
            # special treatment for embedding layer
            if (not with_shared) and "shared_layers" in layer_name:
                continue
            file_name = f"./tmp_layer_files/{layer_name}.tmp"
            paddle.save(layer, file_name)
            ans.append((layer_name, file_name))
            print(f"save layer {layer_name} to {file_name}")
        return ans

    def sort_layers(self, layers: list):
        def priority(elem):
            layer_name = elem[0]
            if "shared_layers" in layer_name:
                return -0.5
            match = re.search(
                r"^_layers((\.\d+)+|(\.shared_layers\.[^\.]+))", layer_name
            )
            assert match, f"{layer_name} not a valid layer name"
            return float(match.group(1).lstrip("."))

        # strictly sort layers
        print("before sort {}".format("|".join([e[0] for e in layers])))
        layers.sort(key=priority)
        # unique
        unique_layers = []
        for e in layers:
            if unique_layers and e[0] == unique_layers[-1][0]:
                continue
            unique_layers.append(e)
        print("after sort {} ".format("|".join([e[0] for e in unique_layers])))
        return unique_layers

    def segment_layers(
        self,
        layers: list,
        config: ParallelConfig,
        segment_method: str = "layer",
    ):
        layer_num = len(layers)
        stage_num = config.pp * config.vpp

        # segment by weights
        def segment_by_layer():
            # assume model is of the structure below
            # embedding -> n*(transformer layer) -> [optional output layer]
            # segment index
            weights = [0 for _ in range(layer_num)]
            non_zero_layers = range(1, layer_num - 1)
            # input layer is embedding
            if self._transformer_layer_num:
                assert self._transformer_layer_num < layer_num
                non_zero_layers = range(1, 1 + self._transformer_layer_num)
            for i in non_zero_layers:
                weights[i] = 1

            part_size = sum(weights) // stage_num
            result = [0 for _ in range(stage_num + 1)]
            memory_counter = 0
            result_idx = 1
            for idx, weight in enumerate(weights):
                memory_counter += weight
                if memory_counter == part_size:
                    result[result_idx] = idx + 1
                    result_idx += 1
                    memory_counter = 0
            result[stage_num] = layer_num
            return result

        def segment_uniform():
            result = [0 for _ in range(stage_num + 1)]
            part_size = math.floor(layer_num / stage_num)
            extra_layers = layer_num % stage_num
            for i in range(1, stage_num):
                offset = 1 if i > (stage_num - extra_layers) else 0
                result[i] = int(
                    min(result[i - 1] + part_size + offset, layer_num)
                )
            result[stage_num] = layer_num
            return result

        result = (
            segment_uniform()
            if (segment_method == "uniform")
            else segment_by_layer()
        )
        index_segments = [[] for _ in range(config.pp)]
        for i in range(stage_num):
            index_segments[i % config.pp].append((result[i], result[i + 1]))

        # name layers
        segments = [[] for i in range(config.pp)]
        for i in range(config.pp):
            for start, end in index_segments[i]:
                for j in range(start, end):
                    if config.vpp > 1:
                        segments[i].append(
                            (
                                [f"_layers.{start}.{j - start}"],
                                layers[j][1],
                            )
                        )
                    else:
                        segments[i].append(([f"_layers.{j}"], layers[j][1]))

        shared_layer_exist = any(
            "_layers.shared_layers" in e[0] for e in layers
        )
        if shared_layer_exist:
            # special treatment for shared layer
            if config.vpp > 1:
                segments[0] = [
                    ([layers[0][0], segments[0][0][0][0]], layers[0][1]),
                    *segments[0][1:],
                ]
            else:
                segments[0] = [([layers[0][0]], layers[0][1]), *segments[0][1:]]

            for i in range(1, config.pp):
                segments[i] = [([layers[0][0]], layers[0][1])] + segments[i]

        for pp_rank, segs in enumerate(segments):
            print(f"segment result for pp_rank {pp_rank}:")
            print(50 * "=")
            for seg in segs:
                print(f"{seg[0]} => {seg[1]}")
        return segments

    def merge_layers(self, layers_segment: list, save_dir: str):
        params = OrderedDict()
        opt = OrderedDict()
        master_weights = OrderedDict()
        renaming_manager = LayerReNamingManager()

        def merge(src, dst, map_k=None):
            for k, v in src.items():
                k = map_k(k) if map_k is not None else k
                dst[k] = v

        lr_scheduler = None
        for layer_names, file_path in layers_segment:
            print(f"load {file_path}")
            layer = paddle.load(file_path)

            def get_param_name_mapper(layer_name):
                # replace layer name
                def map_param_name(param_name):
                    layer_pre = self._extract_layer_name(param_name)
                    return layer_name + param_name[len(layer_pre) :]

                return map_param_name

            (
                layer_params,
                layer_opt,
                layer_master_weight,
            ) = self._map_tensor_names(
                layer["params"],
                layer["opt"],
                layer["master_weights"],
                renaming_manager,
            )
            for layer_name in layer_names:
                merge(layer_params, params, get_param_name_mapper(layer_name))
            merge(layer_opt, opt)
            merge(layer_master_weight, master_weights)
            lr_scheduler = layer["LR_Scheduler"]

        opt = self._pack_opt_state_dict(opt, master_weights, lr_scheduler)
        paddle.save(params, save_dir + "/model.pdparams")
        paddle.save(opt, save_dir + "/model_state.pdopt")

    def _pack_opt_state_dict(self, opt, master_weights, lr_scheduler):
        opt["master_weights"] = master_weights
        opt["LR_Scheduler"] = lr_scheduler
        return opt

    def _extract_layer_name(self, param_name: str):
        match = re.search(
            r"^_layers((\.\d+)+|(\.shared_layers\.[^\.]+))", param_name
        )
        layer_name = ""
        return "" if (not match) else match.group()

    # map opt names to tensor name
    def _opt_name_to_tname(self, tensor_names, opt_names):
        tensor_names = set(tensor_names)
        all_names = []
        all_names.extend(list(tensor_names))
        all_names.extend(opt_names)
        all_names.sort()
        pre_t_name = ""
        opt_to_t = {}
        for n in all_names:
            if n in tensor_names:
                # we get a param
                pre_t_name = n
            else:
                assert pre_t_name
                opt_to_t[n] = pre_t_name
        return opt_to_t

    def _map_tensor_names(self, params, opt, master_weights, renaming_manager):
        opt_renamed = OrderedDict()
        master_weights_renamed = OrderedDict()
        # old name to new name
        t_name_mapping = {}
        # map tensor names
        for k, v in params.items():
            t_name_mapping[v.name] = renaming_manager.get_new_param_name(v.name)
            v.name = t_name_mapping[v.name]
        # map opt names
        opt_to_tname = self._opt_name_to_tname(
            t_name_mapping.keys(), opt.keys()
        )
        for k, v in opt.items():
            old_t_name = opt_to_tname[k]
            t_name = t_name_mapping[old_t_name]
            opt_name = t_name + k[len(old_t_name) :]
            v.name = opt_name
            opt_renamed[opt_name] = v

        # map master names
        for k, v in master_weights.items():
            t_name = t_name_mapping[k]
            v.name = t_name + v.name[len(k) :]
            master_weights_renamed[t_name] = v
        return (params, opt_renamed, master_weights_renamed)


def parse_args():
    parser = argparse.ArgumentParser(
        prog='model converter', description='converter a model'
    )
    parser.add_argument(
        '--src_path',
        type=str,
        default="./output/epoch_0_step_30",
        help='path of the model to convert',
    )

    parser.add_argument(
        '--dst_path',
        type=str,
        default="./test_adapt",
        help='path to saved the converted model',
    )

    parser.add_argument(
        '--src_mp',
        type=int,
        default=2,
        help='mp degree of the origin training task that dumped this model',
    )

    parser.add_argument(
        '--src_pp',
        type=int,
        default=2,
        help='pp degree of the origin training task that dumped this model',
    )

    parser.add_argument(
        '--src_vp',
        type=int,
        default=2,
        help='vp degree of the origin training task that dumped this model',
    )

    parser.add_argument(
        '--dst_mp',
        type=int,
        default=None,
        help='mp degree of the origin training task that dumped this model',
    )

    parser.add_argument(
        '--dst_pp',
        type=int,
        default=None,
        help='pp degree of the expected training task that would recover this model',
    )

    parser.add_argument(
        '--dst_vp',
        type=int,
        default=2,
        help='vp degree of the expected training task that would recover this model',
    )

    parser.add_argument(
        '--sharding',
        type=int,
        default=1,
        help=" sharding degree of both the origin training task that dumped this model and the expected training task that would recover this model",
    )

    parser.add_argument(
        '--method',
        type=str,
        default="adapt_model",
        help='vp degree of the expected training task that would recover this model',
    )

    parser.add_argument(
        '--segment_method',
        type=str,
        default="layer",
        help='method to segment layers to pp or vp stages',
    )

    parser.add_argument(
        '--transformer_layer_num',
        type=int,
        default=0,
        help='transformer_layer_num of the model',
    )
    # assume model is of the structure below
    # embedding -> n*[transformer layer] -> optional output layer

    args = parser.parse_args()

    if args.dst_mp is None:
        args.dst_mp = args.src_mp
    if args.dst_pp is None:
        args.dst_pp = args.src_pp

    assert args.src_mp == args.dst_mp, (
        f"src mp {args.src_mp} dst mp {args.dst_mp}"
    )

    assert args.method in [
        'peek_model',
        'adapt_model',
    ], "method should be in ['peek_model', 'adapt_model']"
    assert args.segment_method in [
        "uniform",
        "layer",
    ], "segment_method should be 'uniform' or 'layer"

    print(
        f"adapt model dumped by task with pp degree:{args.src_pp}, vp degree:{args.src_vp}, mp degree:{args.src_mp} to task with pp degree:{args.dst_pp}, vp degree:{args.dst_vp}, mp degree:{args.dst_mp}"
    )

    return args


def adaptor_from_args(args):
    src_parallel_config = ParallelConfig(
        args.src_mp, args.src_pp, args.src_vp, args.sharding
    )

    dst_parallel_config = ParallelConfig(
        args.dst_mp, args.dst_pp, args.dst_vp, args.sharding
    )

    adaptor = PipeLineModelAdaptor(
        src_parallel_config,
        dst_parallel_config,
        args.transformer_layer_num,
        args.segment_method,
    )
    return adaptor


def main():
    args = parse_args()
    adaptor = adaptor_from_args(args)
    if args.method == "peek_model":
        adaptor.peek_model(args.dst_path)
    elif args.method == "adapt_model":
        adaptor.apply(args.src_path, args.dst_path)


if __name__ == "__main__":
    """
    Usage:
        python pp_parallel_adaptor.py --src_mp xxx --src_path xxx --method \
          adapt_model/peek_model --dst_path xxx --sharding xxx --segment_method xxx --transformer_layer_num xxx

    for the meaning of a specific arg, please use:
        python  pp_parallel_adaptor.py  -h
    """
    main()