# Copyright (c) 2025 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 numpy as np
import paddle
import paddle.distributed as dist
from paddle.distributed import fleet
from paddle.distributed.fleet.base import topology
from paddle.distributed.fleet.base.topology import (
    CommunicateTopology,
    HybridCommunicateGroup,
)
from paddle.distributed.fleet.layers.mpu.random import get_rng_state_tracker

from paddlenlp.transformers.model_utils import unwrap_model
from paddlenlp.utils.log import logger


class ReshardController:
    def __init__(
        self,
        tensor_parallel_degree,
        pipeline_parallel_degree=1,
        sharding_parallel_degree=1,
        sep_parallel_degree=1,
        seed=100,
    ):
        self.tensor_parallel_degree = tensor_parallel_degree
        self.pipeline_parallel_degree = pipeline_parallel_degree
        self.sharding_parallel_degree = sharding_parallel_degree
        self.sep_parallel_degree = sep_parallel_degree
        self.seed = seed
        self.orig_rng_state = paddle.get_rng_state()
        self.orig_cuda_rng_state = self._get_rng_state()
        self.train_hcg = fleet.get_hybrid_communicate_group()
        self.train_tp_group, self.train_dp_group, self.train_sdp_group = (
            self.train_hcg.get_model_parallel_group(),
            self.train_hcg.get_data_parallel_group(),
            self.train_hcg.get_sharding_parallel_group(),
        )
        self.infer_tp_group, self.infer_dp_group, self.infer_sdp_group = self.init_rollout_env()
        self.set_train_env()

        self.is_train = True

    def init_rollout_env(self):
        world_size = dist.get_world_size()
        infer_topo = CommunicateTopology(
            hybrid_group_names=["data", "pipe", "sharding", "sep", "model"],
            dims=[
                world_size
                // self.tensor_parallel_degree
                // self.pipeline_parallel_degree
                // self.sharding_parallel_degree
                // self.sep_parallel_degree,
                self.pipeline_parallel_degree,
                self.sharding_parallel_degree,
                self.sep_parallel_degree,
                self.tensor_parallel_degree,
            ],
        )
        infer_hcg = HybridCommunicateGroup(infer_topo)
        infer_tp_group = infer_hcg.get_model_parallel_group()
        infer_dp_group = infer_hcg.get_data_parallel_group()
        infer_sdp_group = infer_hcg.get_sharding_parallel_group()
        return (infer_tp_group, infer_dp_group, infer_sdp_group)

    def _get_rng_state(self):
        """get_rng_state"""
        origin_rng_state = paddle.get_cuda_rng_state()
        paddle.seed(self.seed)
        rng_state = paddle.get_cuda_rng_state()
        paddle.set_cuda_rng_state(origin_rng_state)
        return rng_state

    def set_rollout_env(self, msg=""):
        if "model_parallel_rng" not in get_rng_state_tracker().states_:
            local_seed = 2025 + 1 + self.train_hcg.get_model_parallel_rank()
            get_rng_state_tracker().add("model_parallel_rng", local_seed)

        paddle.set_rng_state(self.orig_cuda_rng_state)

        hcg = fleet.get_hybrid_communicate_group()
        hcg.get_model_parallel_group = lambda: self.infer_tp_group
        hcg.get_model_parallel_world_size = lambda: self.infer_tp_group.nranks
        hcg.get_model_parallel_rank = lambda: self.infer_tp_group.rank
        hcg.get_sharding_parallel_group = lambda: self.infer_sdp_group
        hcg.get_data_parallel_group = lambda: self.infer_dp_group
        topology._HYBRID_PARALLEL_GROUP.get_model_parallel_group = lambda: self.infer_tp_group
        topology._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size = lambda: self.infer_tp_group.nranks
        topology._HYBRID_PARALLEL_GROUP.get_model_parallel_rank = lambda: self.infer_tp_group.rank
        self.log(msg, False)
        # self.is_train = False

    def set_train_env(self, msg=""):
        hcg = fleet.get_hybrid_communicate_group()
        hcg.get_model_parallel_group = lambda: self.train_tp_group
        hcg.get_model_parallel_world_size = lambda: self.train_tp_group.nranks
        hcg.get_model_parallel_rank = lambda: self.train_tp_group.rank
        hcg.get_sharding_parallel_group = lambda: self.train_sdp_group
        hcg.get_data_parallel_group = lambda: self.train_dp_group
        topology._HYBRID_PARALLEL_GROUP.get_model_parallel_group = lambda: self.train_tp_group
        topology._HYBRID_PARALLEL_GROUP.get_model_parallel_world_size = lambda: self.train_tp_group.nranks
        topology._HYBRID_PARALLEL_GROUP.get_model_parallel_rank = lambda: self.train_tp_group.rank

        paddle.set_rng_state(self.orig_rng_state)
        self.log(msg, True)
        # self.is_train = True

    def log(self, msg, is_train=False):
        msg = f"for {msg}" if len(msg) > 0 else ""
        if is_train:
            logger.warning(
                f"Recover train env done {msg}. [Global TP]: {fleet.get_hybrid_communicate_group().get_model_parallel_world_size()}, [Train TP]: {self.train_tp_group.nranks}, [Infer TP]: {self.infer_tp_group.nranks}"
            )
        else:
            logger.warning(
                f"Set rollout env done {msg}. [Global TP]: {fleet.get_hybrid_communicate_group().get_model_parallel_world_size()}, [Train TP]: {self.train_tp_group.nranks}"
            )


@paddle.no_grad()
def pp_reshard(tgt_tensor, src_model_state_dict, src_tensor_meta_info, pp_rank, pp_group):
    src_tensor_key = src_tensor_meta_info["pipeline_key"]
    src_tensor_pp_rank = src_tensor_meta_info["pipeline_src_rank"]
    src_tensor_shape = src_tensor_meta_info["shape"]

    if src_tensor_pp_rank == pp_rank:
        src_tensor = src_model_state_dict.pop(src_tensor_key)
        resharded_tensor = src_tensor.clone()
        cpu_src_tensor = src_tensor.pin_memory()
        cpu_src_tensor._share_buffer_to(src_tensor)
    else:
        resharded_tensor = paddle.empty(src_tensor_shape)

    resharded_tensor = resharded_tensor.astype(tgt_tensor.dtype)
    dist.broadcast(resharded_tensor, src=pp_group.ranks[src_tensor_pp_rank], group=pp_group, sync_op=True)
    return resharded_tensor


@paddle.no_grad()
def mp_reshard(
    src_tensor,
    tgt_tensor,
    meta_dict,
    train_tp_group,
    rollout_tp_group,
):
    if rollout_tp_group.nranks == train_tp_group.nranks:
        return src_tensor

    if meta_dict["is_distributed"]:
        res = []
        if train_tp_group.nranks > 1:
            paddle.distributed.all_gather(res, src_tensor, group=train_tp_group, sync_op=True)
        else:
            res = [src_tensor]
        if hasattr(tgt_tensor, "is_distributed") and tgt_tensor.is_distributed:
            assert hasattr(tgt_tensor, "split_axis"), f"{tgt_tensor.name} has no split_axis!"
            concat_tensor = paddle.concat(res, meta_dict["split_axis"])
            del res
            all_parts = paddle.split(concat_tensor, rollout_tp_group.nranks, tgt_tensor.split_axis)
            del concat_tensor
            return all_parts[rollout_tp_group.rank]
        else:
            return paddle.concat(res, meta_dict["split_axis"])
    return src_tensor


def init_reshard_mappings(model, training_args, pp_rank, pp_group):
    global_meta_dict = {}
    if training_args.pipeline_parallel_degree > 1:
        model._layers._set_pipeline_name_mapping()
        local_name_mapping_dict = model._layers._single_to_pp_mapping
    else:
        local_name_mapping_dict = {}
        for k in model.state_dict():
            local_name_mapping_dict[k] = k.replace("_layers.", "")
    local_model_state_dict = unwrap_model(model).state_dict()
    local_meta_dict = {}
    for k, v in local_name_mapping_dict.items():
        if training_args.pipeline_parallel_degree == 1:
            k = k.replace("_layers.", "")
        pipeline_key = v
        pipeline_tensor = local_model_state_dict[pipeline_key]
        local_meta_dict[k] = {
            "pipeline_key": pipeline_key,
            "pipeline_src_rank": pp_rank,
            "shape": pipeline_tensor.shape,
        }
        local_meta_dict[k]["is_distributed"] = False
        if hasattr(pipeline_tensor, "is_distributed"):
            local_meta_dict[k]["is_distributed"] = pipeline_tensor.is_distributed
        local_meta_dict[k]["split_axis"] = None
        if hasattr(pipeline_tensor, "split_axis"):
            local_meta_dict[k]["split_axis"] = pipeline_tensor.split_axis
    if training_args.pipeline_parallel_degree > 1:
        gathered_local_meta_dict = []
        dist.all_gather_object(gathered_local_meta_dict, local_meta_dict, group=pp_group)
    else:
        gathered_local_meta_dict = [local_meta_dict]
    for meta_dict in gathered_local_meta_dict:
        global_meta_dict.update(meta_dict)
    return global_meta_dict


@paddle.no_grad()
def reshard_to_rollout(
    train_model, rollout_model, global_meta_dict, pp_rank, pp_group, rollout_tp_group, train_tp_group
):
    train_model_state_dict = train_model.state_dict()
    rollout_model_state_dict = rollout_model.state_dict()
    param_numel = [(k, np.prod(v.shape)) for k, v in rollout_model_state_dict.items()]
    param_numel.sort(key=lambda x: x[1], reverse=True)

    for k, _ in param_numel:
        v = rollout_model_state_dict[k]
        resharded_tensor = pp_reshard(v, train_model_state_dict, global_meta_dict[k], pp_rank, pp_group)
        resharded_tensor = mp_reshard(
            resharded_tensor,
            v,
            global_meta_dict[k],
            train_tp_group,
            rollout_tp_group,
        )
        assert resharded_tensor.dtype == v.dtype, f"dtype wrong {k} {resharded_tensor.dtype} {v.dtype}"
        assert resharded_tensor.shape == v.shape, f"shape wrong {k} {resharded_tensor.shape} {v.shape}"
        resharded_tensor._share_buffer_to(v)
        resharded_tensor._clear()

    missing_keys = train_model_state_dict.keys()
    num_missing_keys = len(missing_keys)
    assert num_missing_keys == 0, f"missing {num_missing_keys} keys after reshard policy: {missing_keys}"
    logger.warning("[Reshard] Done")
