# Copyright (c) 2022 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 collections
import logging
import os
import warnings
from functools import reduce

from paddle.base.framework import generate_control_dev_var_name
from paddle.distributed.io import is_persistable
from paddle.framework import core

# logging.basicConfig(
#    format='%(levelname)s - %(asctime)s - %(pathname)s: %(lineno)s - %(message)s', level=logging.INFO)
# logger = logging.getLogger(__name__)

OP_NAME_SCOPE = "op_namescope"
CLIP_OP_NAME_SCOPE = "gradient_clip"
STEP_COUNTER = "@PS_STEP_COUNTER@"
LEARNING_RATE_DECAY_COUNTER = "@LR_DECAY_COUNTER@"

OP_ROLE_VAR_ATTR_NAME = core.op_proto_and_checker_maker.kOpRoleVarAttrName()
RPC_OP_ROLE_ATTR_NAME = core.op_proto_and_checker_maker.kOpRoleAttrName()
RPC_OP_ROLE_ATTR_VALUE = core.op_proto_and_checker_maker.OpRole.RPC
op_role = core.op_proto_and_checker_maker.OpRole
op_role_attr_name = core.op_proto_and_checker_maker.kOpRoleAttrName()
LR_SCHED_OP_ROLE_ATTR_VALUE = core.op_proto_and_checker_maker.OpRole.LRSched
OPT_OP_ROLE_ATTR_VALUE = core.op_proto_and_checker_maker.OpRole.Optimize
backward = core.op_proto_and_checker_maker.OpRole.Backward
OP_DEVICE_KEY = core.op_proto_and_checker_maker.kOpDeviceAttrName()

DEVICE_LIST = ["cpu", "gpu", "xpu"]
COMMUNICATE_OPS_TYPE = ["send", "recv", "fetch_barrier", "send_barrier"]
SPARSE_OP_LIST = ["lookup_table", "lookup_table_v2"]
SPARSE_OP_TYPE_DICT = {"lookup_table": "W", "lookup_table_v2": "W"}
SPARSE_GRAD_OP_TYPE_DICT = {
    "lookup_table_grad": "W",
    "lookup_table_v2_grad": "W",
}
DEFAULT_DEVICE = 'cpu'

DATA_NORM_NAME = [".batch_size", ".batch_sum", ".batch_square_sum"]
DATA_NORM_GRAD_NAME = [x + "@GRAD" for x in DATA_NORM_NAME]


def logger_config(log_path, logging_name):
    logger = logging.getLogger(logging_name)
    logger.setLevel(level=logging.WARNING)
    handler = logging.FileHandler(
        log_path, mode='a', encoding='UTF-8', delay=True
    )
    handler.setLevel(logging.INFO)
    formatter = logging.Formatter(
        '%(levelname)s - %(asctime)s - %(pathname)s: %(lineno)s - %(message)s'
    )
    handler.setFormatter(formatter)
    console = logging.StreamHandler()
    console.setLevel(logging.DEBUG)
    logger.addHandler(handler)
    logger.addHandler(console)
    return logger


ps_log_root_dir = './ps_log/'
logger = logger_config(
    log_path='./ps_usr_print_log', logging_name='ps_usr_print_log'
)


class DistributedMode:
    SYNC = 0
    ASYNC = 1
    HALF_ASYNC = 2
    GEO = 3
    FL = 4
    NU = 5


class TrainerRuntimeConfig:
    def __init__(self, valid_strategy):
        self.mode = None
        num_threads = os.getenv("CPU_NUM", "1")
        send_queue_size = num_threads
        k_steps = valid_strategy.a_sync_configs["k_steps"]

        if not valid_strategy.a_sync and k_steps == 0:
            self.mode = DistributedMode.SYNC

        if valid_strategy.a_sync and k_steps == 0:
            self.mode = DistributedMode.ASYNC

        if valid_strategy.a_sync and k_steps > 0:
            self.mode = DistributedMode.GEO
            send_queue_size = k_steps

        self.runtime_configs = {}
        self.runtime_configs['communicator_max_merge_var_num'] = os.getenv(
            "FLAGS_communicator_max_merge_var_num", send_queue_size
        )
        self.runtime_configs['communicator_send_queue_size'] = os.getenv(
            "FLAGS_communicator_send_queue_size", send_queue_size
        )
        self.runtime_configs['communicator_independent_recv_thread'] = (
            os.getenv("FLAGS_communicator_independent_recv_thread", "1")
        )
        self.runtime_configs['communicator_min_send_grad_num_before_recv'] = (
            os.getenv(
                "FLAGS_communicator_min_send_grad_num_before_recv", num_threads
            )
        )
        self.runtime_configs['communicator_thread_pool_size'] = os.getenv(
            "FLAGS_communicator_thread_pool_size", "5"
        )
        self.runtime_configs['communicator_send_wait_times'] = os.getenv(
            "FLAGS_communicator_send_wait_times", "5"
        )
        self.runtime_configs['communicator_is_sgd_optimizer'] = os.getenv(
            "FLAGS_communicator_is_sgd_optimizer", "1"
        )

    def get_communicator_flags(self):
        need_keys = []
        num_threads = os.getenv("CPU_NUM", "1")
        mode_str = ""
        if self.mode is None or self.mode == DistributedMode.ASYNC:
            need_keys = self.runtime_configs.keys()
            mode_str = "async"
        elif (
            self.mode == DistributedMode.SYNC
            or self.mode == DistributedMode.HALF_ASYNC
        ):
            mode_str = "sync or half_async"
            need_keys = [
                'communicator_max_merge_var_num',
                'communicator_send_wait_times',
                'communicator_thread_pool_size',
                'communicator_send_queue_size',
            ]
        elif self.mode == DistributedMode.GEO:
            mode_str = "GEO"
            need_keys = [
                'communicator_thread_pool_size',
                'communicator_send_wait_times',
                'communicator_max_merge_var_num',
                'communicator_send_queue_size',
            ]
        else:
            raise ValueError("Unsupported Mode")

        if (
            self.mode == DistributedMode.SYNC
            or self.mode == DistributedMode.HALF_ASYNC
        ):
            max_merge_var_num = self.runtime_configs[
                'communicator_max_merge_var_num'
            ]
            send_queue_size = self.runtime_configs[
                'communicator_send_queue_size'
            ]
            if max_merge_var_num != num_threads:
                print(
                    f'WARNING: In {mode_str} mode, communicator_max_merge_var_num '
                    'must be equal to CPU_NUM. But received, '
                    f'communicator_max_merge_var_num = {max_merge_var_num}, CPU_NUM = '
                    f'{num_threads}. communicator_max_merge_var_num will be forced to {num_threads}.'
                )
                self.runtime_configs['communicator_max_merge_var_num'] = (
                    num_threads
                )
            if send_queue_size != num_threads:
                print(
                    f'WARNING: In {mode_str} mode, communicator_send_queue_size '
                    'must be equal to CPU_NUM. But received, '
                    f'communicator_send_queue_size = {send_queue_size}, CPU_NUM = '
                    f'{num_threads}. communicator_send_queue_size will be forced to {num_threads}.'
                )
                self.runtime_configs['communicator_send_queue_size'] = (
                    num_threads
                )

        return {key: str(self.runtime_configs[key]) for key in need_keys}


def get_lr_ops(program):
    lr_ops = []
    for index, op in enumerate(program.global_block().ops):
        role_id = int(op.attr(RPC_OP_ROLE_ATTR_NAME))
        if role_id == int(LR_SCHED_OP_ROLE_ATTR_VALUE) or role_id == int(
            LR_SCHED_OP_ROLE_ATTR_VALUE
        ) | int(OPT_OP_ROLE_ATTR_VALUE):
            lr_ops.append(op)
    return lr_ops


def get_optimize_ops(_program, remote_sparse=[]):
    block = _program.global_block()
    opt_ops = []
    for op in block.ops:
        if _is_opt_role_op(op):
            if (
                len(remote_sparse) > 0
                and op.input("Param")[0] not in remote_sparse
            ):  # for fl: only delete remote sparse optimize
                continue
            # delete clip op from opt_ops when run in Parameter Server mode
            if (
                OP_NAME_SCOPE in op.all_attrs()
                and CLIP_OP_NAME_SCOPE in op.attr(OP_NAME_SCOPE)
            ):
                op._set_attr(
                    "op_role",
                    int(core.op_proto_and_checker_maker.OpRole.Backward),
                )
                continue
            opt_ops.append(op)
    return opt_ops


def get_datanorm_ops(_program):
    block = _program.global_block()
    opt_ops = []
    for op in block.ops:
        if op.type == 'data_norm':
            opt_ops.append(op)
    return opt_ops


def get_dist_env():
    trainer_id = int(os.getenv('PADDLE_TRAINER_ID', '0'))
    trainer_endpoints = ''
    current_endpoint = ''
    num_trainers = 0
    if os.getenv('PADDLE_TRAINER_ENDPOINTS'):
        trainer_endpoints = os.getenv('PADDLE_TRAINER_ENDPOINTS')
        current_endpoint = trainer_endpoints.split(',')[trainer_id]
        num_trainers = len(trainer_endpoints.split(','))

    return {
        'trainer_id': trainer_id,
        'num_trainers': num_trainers,
        'current_endpoint': current_endpoint,
        'trainer_endpoints': trainer_endpoints,
    }


def get_role_id(role_maker):
    try:
        return role_maker._role_id()
    except Exception:
        return role_maker.role_id()


def get_ps_endpoint(role_maker):
    try:
        return role_maker._get_pserver_endpoints()[get_role_id(role_maker)]
    except Exception:
        return role_maker.get_pserver_endpoints()[get_role_id(role_maker)]


def get_ps_endpoints(role_maker):
    try:
        return role_maker._get_pserver_endpoints()
    except Exception:
        return role_maker.get_pserver_endpoints()


def get_heter_worker_endpoint(role_maker):
    return role_maker._get_heter_worker_endpoint()


def get_trainer_endpoint(role_maker):
    return role_maker._get_trainer_endpoint()


def get_trainer_endpoints(role_maker):
    return role_maker._get_trainer_endpoints()


def get_previous_stage_trainers(role_maker):
    try:
        return role_maker._get_previous_trainers()
    except Exception:
        return role_maker.get_previous_trainers()


def is_distributed_sparse_op(op):
    if op.type in SPARSE_OP_LIST and op.attr('is_distributed') is True:
        return True

    if (
        op.type == "distributed_lookup_table"
        and op.attr('is_distributed') is True
    ):
        return True

    return False


def get_sparse_tablename(op):
    return op.input("W")[0]


def is_sparse_op(op):
    if (
        op.type in SPARSE_OP_LIST
        and op.attr('is_sparse') is True
        and op.attr('is_distributed') is False
    ):
        return True

    if (
        op.type == "distributed_lookup_table"
        and op.attr('is_distributed') is False
    ):
        return True

    return False


def get_sparse_tablenames(programs, is_distributed):
    tablenames = set()
    for program in programs:
        if is_distributed:
            for op in program.global_block().ops:
                if is_distributed_sparse_op(op):
                    tablenames.add(get_sparse_tablename(op))
        else:
            for op in program.global_block().ops:
                if is_sparse_op(op):
                    tablenames.add(get_sparse_tablename(op))
    return list(tablenames)


def get_trainers(role_maker):
    try:
        return role_maker._worker_num()
    except Exception:
        return role_maker.worker_num()


def get_dense_send_context(
    program,
    send_ctx,
    idx,
    merged_dense_pairs,
    trainer_id,
    split_dense_table=False,
):
    if len(merged_dense_pairs) < 1:
        return idx
    if not split_dense_table:
        dense_pairs = []
        data_norm_pairs = []
        for merged in merged_dense_pairs:
            is_data_norm = False
            grad = merged[1]
            varname = grad.merged_var.name
            for name in DATA_NORM_GRAD_NAME:
                if varname.endswith(name):
                    is_data_norm = True
            if is_data_norm:
                data_norm_pairs.append(merged)
            else:
                dense_pairs.append(merged)

        # simple dense table
        origin_varnames = []
        var_numel = 0
        for merged in dense_pairs:
            grad = merged[1]
            origin_varnames.append(grad.merged_var.name)
            var = program.global_block().vars[grad.merged_var.name]
            var_numel += reduce(lambda x, y: x * y, var.shape, 1)
        grad_name = "Dense@GRAD_" + str(idx)
        aggregate = True
        # print("public get_dense_send_context dense_table:", grad_name,
        #      var_numel, origin_varnames)
        from paddle.base.core import CommContext

        dense_ctx = CommContext(
            grad_name,
            [grad_name],
            ["127.0.0.1:6071"],
            [var_numel],
            origin_varnames,
            trainer_id,
            aggregate,
            False,
            False,
            idx,
            False,
            False,
            id(program),
            [],
        )
        send_ctx[grad_name] = dense_ctx
        idx += 1

        if len(data_norm_pairs) <= 0:
            return idx

        # data norm table
        origin_varnames = []
        var_numel = 0
        for merged in data_norm_pairs:
            grad = merged[1]
            origin_varnames.append(grad.merged_var.name)
            var = program.global_block().vars[grad.merged_var.name]
            var_numel += reduce(lambda x, y: x * y, var.shape, 1)
        grad_name = "DataNorm@GRAD_" + str(idx)
        aggregate = True
        # print("public get_dense_send_context data_norm table:", grad_name,
        #      var_numel, origin_varnames)
        from paddle.base.core import CommContext

        data_norm_ctx = CommContext(
            grad_name,
            [grad_name],
            ["127.0.0.1:6071"],
            [var_numel],
            origin_varnames,
            trainer_id,
            aggregate,
            False,
            False,
            idx,
            False,
            True,
            id(program),
            [],
        )
        send_ctx[grad_name] = data_norm_ctx
        idx += 1
    else:
        for merged in merged_dense_pairs:
            grad = merged[1]
            origin_varname = grad.merged_var.name
            var = program.global_block().vars[origin_varname]
            var_numel = reduce(lambda x, y: x * y, var.shape, 1)
            grad_name = origin_varname
            aggregate = True
            from paddle.base.core import CommContext

            dense_ctx = CommContext(
                grad_name,
                [grad_name],
                ["127.0.0.1:6071"],
                [var_numel],
                [origin_varname],
                trainer_id,
                aggregate,
                False,
                False,
                idx,
                False,
                False,
                id(program),
                [],
            )
            send_ctx[grad_name] = dense_ctx
            idx += 1
    return idx


def get_geo_trainer_send_context(attrs):
    if attrs['ps_mode'] != DistributedMode.GEO:
        raise ValueError(
            f"ps mode: {attrs['ps_mode']} not matched get_geo_trainer_send_context",
        )
    send_ctx = {}
    trainer_id = get_role_id(attrs['role_maker'])
    origin_programs = attrs['origin_main_programs']
    idx = 0  # table idx

    distributed_varnames = get_sparse_tablenames(origin_programs, True)
    for i, program in enumerate(origin_programs):
        merged_sparse_pairs = attrs['merged_sparse_pairs'][i]
        for merged in merged_sparse_pairs:
            param, grad = merged
            grad_name = grad.merged_var.name
            param_name = param.merged_var.name
            if param_name in attrs['remote_sparse']:  # for recall/ncf model
                continue

            is_distributed = (
                True if param_name in distributed_varnames else False
            )
            var = program.global_block().vars[grad.merged_var.name]
            var_numel = reduce(lambda x, y: x * y, var.shape[1:], 1)
            from paddle.base.core import CommContext

            print(
                "public get_the_geo_send_context sparse: ", grad_name, var_numel
            )
            sparse_ctx = CommContext(
                grad_name,
                [grad_name],
                ["127.0.0.1:6071"],
                [var_numel],
                [grad_name],
                trainer_id,
                True,
                True,
                is_distributed,
                idx,
                False,
                False,
                id(program),
                [],
            )
            idx += 1
            send_ctx[sparse_ctx.var_name()] = sparse_ctx

    if len(send_ctx) == 0:
        raise ValueError("GeoSGD require sparse parameters in your net.")

    if len(attrs['tensor_table']) > 0 and attrs['is_worker']:
        name, ctx = _step_ctx(idx, attrs['role_maker'])
        send_ctx[name] = ctx

    return send_ctx


def _step_ctx(idx, role_maker):
    name = STEP_COUNTER
    trainer_id = get_role_id(role_maker)
    endpoints = get_ps_endpoints(role_maker)
    sections = [1] * len(endpoints)
    names = [name] * len(endpoints)
    from paddle.base.core import CommContext

    ctx = CommContext(
        name,
        names,
        endpoints,
        sections,
        [name],
        trainer_id,
        True,
        False,
        False,
        idx,
        True,
        False,
        -1,
        [],
    )
    return name, ctx


def get_the_one_send_context(attrs, split_dense_table=False, ep_list=None):
    if ep_list is None:
        ep_list = ["127.0.0.1:6071"]
    send_ctx = {}
    trainer_id = get_role_id(attrs['role_maker'])
    origin_programs = attrs['origin_main_programs']
    print(f"is_heter_ps_mode? {split_dense_table}")

    idx = 0
    distributed_varnames = get_sparse_tablenames(origin_programs, True)
    # print("public distributed_varnames:", distributed_varnames)
    for i, program in enumerate(origin_programs):
        merged_sparse_pairs = attrs['merged_sparse_pairs'][i]
        for merged in merged_sparse_pairs:
            param, grad = merged
            grad_name = grad.merged_var.name
            param_name = param.merged_var.name

            remote_sparse_ids = []
            if param_name in attrs['remote_sparse']:  # for recall/ncf model
                remote_sparse_ids.append(idx)

            splited_varname = []
            for i in range(len(ep_list)):
                splited_varname.append(f"{param_name}.block{i}")

            is_distributed = (
                True if param_name in distributed_varnames else False
            )

            var = program.global_block().vars[grad.merged_var.name]

            shape = list(var.shape)
            shape[0] = 0 if is_distributed else shape[0]

            if grad_name in send_ctx:
                continue
            from paddle.base.core import CommContext

            print(
                "public get_the_one_send_context sparse: ",
                grad_name,
                splited_varname,
                shape,
            )
            sparse_ctx = CommContext(
                grad_name,
                splited_varname,
                ep_list,
                shape,
                [grad_name],
                trainer_id,
                True,
                True,
                is_distributed,
                idx,
                False,
                False,
                id(program),
                remote_sparse_ids,
            )

            idx += 1
            send_ctx[sparse_ctx.var_name()] = sparse_ctx

    for i, program in enumerate(origin_programs):
        merged_dense_pairs = attrs['merged_dense_pairs'][i]
        idx = get_dense_send_context(
            program,
            send_ctx,
            idx,
            merged_dense_pairs,
            trainer_id,
            split_dense_table,
        )

    if len(attrs['tensor_table']) > 0 and attrs['is_worker']:
        name, ctx = _step_ctx(idx, attrs['role_maker'])
        send_ctx[name] = ctx

    return send_ctx


def find_heter_ops(program, default_device="cpu"):
    if default_device not in DEVICE_LIST:
        raise ValueError(
            f"Given device {default_device} is not in device list {DEVICE_LIST}"
        )

    def _is_heter_op(op, current_heter_device, default_device="cpu"):
        heter_devices = list(DEVICE_LIST)
        heter_devices.remove(default_device)
        op_device = op.attr("op_device")
        op_type = op.type
        if op_device in heter_devices:
            return True
        elif (
            op_type in COMMUNICATE_OPS_TYPE
            and current_heter_device != default_device
        ):
            # for distributed communicate ops: send & recv & barrier etc.
            # Todo: need update this method
            # op._set_attr('op_device', current_heter_device)
            return True
        elif op_device is None or op_device == default_device:
            op._set_attr('op_device', default_device)
            return False
        return False

    def _is_same_device(op, pre_device, default_device="cpu"):
        op_device = op.attr("op_device")
        if op_device == pre_device:
            return True
        if pre_device == default_device:
            return True
        return False

    def _append_heter_op(op, current_heter_block_ops, heter_ops):
        op_device = op.attr("op_device")
        if op_device not in heter_ops:
            heter_ops[op_device] = {}
        current_heter_block_ops.append(op)

    origin_program = program.clone()
    block = program.global_block()
    '''
       re-place sum op to fix bug for union forward backward op
    '''
    var2idx = {}
    op_list = list(block.ops)
    op_size = len(op_list)

    for i in range(op_size - 1, -1, -1):
        op_list = list(block.ops)
        op = op_list[i]
        if "_grad" in op.type:
            forward_op_type = op.type.split("_grad")[0]
            if (
                forward_op_type in SPARSE_OP_TYPE_DICT.keys()
                and op.attr('remote_prefetch') is True
            ):
                param_name = op.input(SPARSE_OP_TYPE_DICT[forward_op_type])[0]
                if param_name in var2idx:
                    # insert sum op & remove sum op from var2idx and origin place
                    op_list = list(block.ops)
                    sum_op = op_list[var2idx[param_name]]
                    sum_op_inputs = {
                        sum_op.input_names[0]: [
                            block.vars[input]
                            for input in sum_op.input_arg_names
                        ]
                    }
                    sum_op_outputs = {
                        sum_op.output_names[0]: [
                            block.vars[output]
                            for output in sum_op.output_arg_names
                        ]
                    }
                    block._insert_op(
                        index=i + 1,
                        type=sum_op.type,
                        inputs=sum_op_inputs,
                        outputs=sum_op_outputs,
                        attrs=sum_op.all_attrs(),
                    )
                    block._remove_op(var2idx[param_name] + 1)
                    var2idx.pop(param_name)
                    for var_ in var2idx:
                        var2idx[var_] += 1
            elif forward_op_type == "elementwise_mul":
                """
                get output varname of pre op

                """
                output_vars_no_grad = []
                for key in op.output_names:
                    for varname in op.output(key):
                        if varname == "@EMPTY@":
                            continue
                        if "lod_tensor_blocking_queue" in varname:
                            continue
                        output_vars_no_grad.append(varname.split("@GRAD")[0])
                for no_grad_var in output_vars_no_grad:
                    if no_grad_var in var2idx:
                        """
                        insert sum op & remove sum op from var2idx and origin place

                        """
                        op_list = list(block.ops)
                        sum_op = op_list[var2idx[no_grad_var]]
                        sum_op_inputs = {
                            sum_op.input_names[0]: [
                                block.vars[input]
                                for input in sum_op.input_arg_names
                            ]
                        }
                        sum_op_outputs = {
                            sum_op.output_names[0]: [
                                block.vars[output]
                                for output in sum_op.output_arg_names
                            ]
                        }
                        block._insert_op(
                            index=i + 1,
                            type=sum_op.type,
                            inputs=sum_op_inputs,
                            outputs=sum_op_outputs,
                            attrs=sum_op.all_attrs(),
                        )
                        block._remove_op(var2idx[no_grad_var] + 1)
                        var2idx.pop(no_grad_var)
                        for var_ in var2idx:
                            var2idx[var_] += 1
        else:
            if op.type == "sum":
                var = op.output("Out")[0]
                if "@GRAD" in var:
                    origin_var = var.split("@GRAD")[0]
                    pre_op = op_list[i - 1]
                    if "_grad" in pre_op.type:
                        forward_op_type = pre_op.type.split("_grad")[0]
                        if (
                            forward_op_type in SPARSE_OP_TYPE_DICT.keys()
                            and pre_op.attr('remote_prefetch') is True
                        ):
                            param_name = pre_op.input(
                                SPARSE_OP_TYPE_DICT[forward_op_type]
                            )[0]
                            if param_name == origin_var and op.attr(
                                "op_device"
                            ) == pre_op.attr("op_device"):
                                continue
                            else:
                                var2idx[origin_var] = i
                        elif forward_op_type == "elementwise_mul":
                            output_vars = []
                            for key in pre_op.output_names:
                                for varname in pre_op.output(key):
                                    if varname == "@EMPTY@":
                                        continue
                                    if "lod_tensor_blocking_queue" in varname:
                                        continue
                                    output_vars.append(varname)
                            input_vars = []
                            for key in op.input_names:
                                for varname in op.input(key):
                                    if varname == "@EMPTY@":
                                        continue
                                    if "lod_tensor_blocking_queue" in varname:
                                        continue
                                    input_vars.append(varname)
                            is_match = False
                            for varname in output_vars:
                                if varname in input_vars:
                                    is_match = True
                                    break
                            if is_match:
                                continue
                            else:
                                var2idx[origin_var] = i
                    else:
                        var2idx[origin_var] = i

    origin_program = program.clone()
    block = program.global_block()

    program_block_ops = []
    default_ops = {default_device: {}}
    heter_ops = {}
    block_index = 0

    current_heter_block_ops = []
    current_default_block_ops = []
    current_heter_device = default_device
    is_heter = False
    for op in block.ops:
        if _is_heter_op(op, current_heter_device, default_device):
            # for gpu/xpu-op
            is_heter = True

            # for cpu-op block append
            if len(current_default_block_ops) > 1:
                default_ops[default_device][block_index] = (
                    current_default_block_ops
                )
                program_block_ops.append(current_default_block_ops)
                current_default_block_ops = []
                block_index += 1

            if _is_same_device(op, current_heter_device, default_device):
                # for gpu-op, gpu-op -> gpu-op,...
                current_heter_device = op.attr("op_device")
                _append_heter_op(op, current_heter_block_ops, heter_ops)
            else:
                # for gpu-op -> xpu-op, ...
                op_device = current_heter_block_ops[0].attr("op_device")
                heter_ops[op_device][block_index] = current_heter_block_ops
                program_block_ops.append(current_heter_block_ops)
                block_index += 1
                current_heter_block_ops = []
                current_heter_device = op.attr("op_device")
                _append_heter_op(op, current_heter_block_ops, heter_ops)

        elif is_heter:
            # for gpu/xpu-op -> cpu-op
            op_device = current_heter_block_ops[0].attr("op_device")
            heter_ops[op_device][block_index] = current_heter_block_ops
            program_block_ops.append(current_heter_block_ops)
            block_index += 1
            current_heter_block_ops = []
            current_heter_device = default_device
            is_heter = False
            current_default_block_ops.append(op)
        else:
            # for cpu-op
            current_default_block_ops.append(op)

    if current_default_block_ops != []:
        default_ops[default_device][block_index] = current_default_block_ops
        program_block_ops.append(current_default_block_ops)

    if current_heter_block_ops != []:
        op_device = current_heter_block_ops[0].attr("op_device")
        heter_ops[op_device][block_index] = current_heter_block_ops
        program_block_ops.append(current_heter_block_ops)

    if len(heter_ops) == 0:
        warnings.warn(
            "No heterogeneous OP was found in your program , "
            " please using static.device_guard() to run OPs on different device."
        )

    total_heter_ops = 0
    heter_blocks = 0
    for device in heter_ops.keys():
        heter_block_dict = heter_ops[device]
        heter_blocks += len(heter_block_dict)
        for _, heter_block in heter_block_dict.items():
            total_heter_ops += len(heter_block)
    print(
        f"There are {len(block.ops)} OPs in your main_program, and contains {total_heter_ops} heter-OPs which is made up of {heter_blocks} heter-blocks."
    )

    return origin_program, heter_ops, default_ops, program_block_ops


def union_forward_gradient_op(program_block_ops_list):
    """
    before analyzing the input & output of each block in program_block_list, we should
    union the forward op and corresponding gradient op to eliminate the unnecessary variable
    transmit
    """
    """
    fix for 2emb model, re-place sum op

    """
    block_length = len(program_block_ops_list)
    union_program_block_ops_list = []
    assert block_length % 2 != 0, (
        "the length of program_block_ops_list should be odd"
    )
    for i in range(0, block_length // 2):
        block_op_list = {"forward": program_block_ops_list[i]}
        block_op_list.update(
            {"backward": program_block_ops_list[block_length - 1 - i]}
        )
        union_program_block_ops_list.append(block_op_list)

    block_op_list = {"forward": [], "backward": []}
    for op in program_block_ops_list[block_length // 2]:
        if "_grad" not in op.type and not (op.type == "sum"):
            block_op_list["forward"].append(op)
        else:
            block_op_list["backward"].append(op)
    union_program_block_ops_list.append(block_op_list)
    return union_program_block_ops_list


def find_block_joints(program, program_block_ops_list, heter_ops):
    block_var_detail = find_entrance_exit_private(
        program, program_block_ops_list
    )
    block_var_detail = entrance_exit_check(
        program, program_block_ops_list, block_var_detail, heter_ops
    )
    block_var_detail = delete_block_useless_exit(
        program, program_block_ops_list, block_var_detail
    )

    return block_var_detail


def find_ops_list_input_output(program, ops_list):
    input_var_list = []
    output_var_list = []
    for op in ops_list:
        inputs = _get_input_map_from_op(program.global_block().vars, op)
        input_var_list += get_varlist_from_op_map(inputs)
        outputs = _get_output_map_from_op(program.global_block().vars, op)
        output_var_list += get_varlist_from_op_map(outputs)

    input_var_list = list(set(input_var_list))
    output_var_list = list(set(output_var_list))
    return input_var_list, output_var_list


def find_entrance_exit_private(program, program_block_ops_list):
    block_var_detail = []
    persistables = []
    for index, block_op_list in enumerate(program_block_ops_list):
        # forward
        block_input, block_output = find_ops_list_input_output(
            program, block_op_list["forward"]
        )
        persistables = screen_persistables(
            program, block_input
        ) + screen_persistables(program, block_output)
        # find entrance & exit
        block_private_vars = list(set(block_input) & set(block_output))
        block_entrance = list(set(block_input) - set(block_private_vars))
        block_exit = list(set(block_output) - set(block_private_vars))
        detail = {
            "forward": {
                "entrance": block_entrance,
                "exit": block_exit,
                "private": block_private_vars,
                "persistables": persistables,
            }
        }

        # backward
        bp_block_input, bp_block_output = find_ops_list_input_output(
            program, block_op_list["backward"]
        )
        bp_persistables = screen_persistables(
            program, bp_block_input
        ) + screen_persistables(program, bp_block_output)
        # find entrance & exit
        bp_block_private_vars = list(set(bp_block_input) & set(bp_block_output))
        bp_block_entrance = list(
            set(bp_block_input) - set(bp_block_private_vars)
        )
        bp_block_exit = list(set(bp_block_output) - set(bp_block_private_vars))
        detail.update(
            {
                "backward": {
                    "entrance": bp_block_entrance,
                    "exit": bp_block_exit,
                    "private": bp_block_private_vars,
                    "persistables": bp_persistables,
                }
            }
        )
        block_var_detail.append(detail)
    return block_var_detail


def entrance_exit_check(
    program, program_block_ops_list, block_var_detail, heter_ops
):
    for index in range(len(block_var_detail) - 1, -1, -1):
        if index - 1 < 0:
            break
        previous_block_exit = block_var_detail[index - 1]["forward"]["exit"]
        previous_block_exit.sort()
        current_block_entrance = block_var_detail[index]["forward"]["entrance"]

        backward_entrance = block_var_detail[index]["backward"]["entrance"]

        forward_all = (
            block_var_detail[index]["forward"]["entrance"]
            + block_var_detail[index]["forward"]["private"]
            + block_var_detail[index]["forward"]["exit"]
        )

        for var in backward_entrance:
            if "@GRAD" not in var and var not in forward_all:
                current_block_entrance.append(var)

        current_block_entrance.sort()

        if previous_block_exit == current_block_entrance:
            continue
        exist_vars = list(
            set(previous_block_exit) & set(current_block_entrance)
        )
        need_add_vars = list(set(current_block_entrance) - set(exist_vars))
        # var in different stage should not be ignored, since they are not placed in the same program & device
        # need_add_vars = find_need_var_from_previous_block(
        #    need_add_vars, block_var_detail, index, heter_ops)

        previous_block_private = block_var_detail[index - 1]["forward"][
            "private"
        ]
        previous_block_entrance = block_var_detail[index - 1]["forward"][
            "entrance"
        ]
        for var in need_add_vars:
            if (
                var not in previous_block_private
                and var not in previous_block_entrance
            ):
                previous_block_entrance.append(var)
            previous_block_exit.append(var)
            if var not in current_block_entrance:
                current_block_entrance.append(var)

    for index in range(0, len(block_var_detail) - 1, 1):
        previous_block_exit = block_var_detail[index + 1]["backward"]["exit"]
        previous_block_exit.sort()
        current_block_entrance = block_var_detail[index]["backward"]["entrance"]

        current_block_entrance.sort()

        if previous_block_exit == current_block_entrance:
            continue
        exist_vars = list(
            set(previous_block_exit) & set(current_block_entrance)
        )
        need_add_vars = list(set(current_block_entrance) - set(exist_vars))
        need_ignore_vars = []
        for var in need_add_vars:
            if "@GRAD" not in var:
                need_ignore_vars.append(var)
        need_add_vars = list(
            set(need_add_vars).difference(set(need_ignore_vars))
        )
        previous_block_private = block_var_detail[index + 1]["backward"][
            "private"
        ]
        previous_block_entrance = block_var_detail[index + 1]["backward"][
            "entrance"
        ]
        for var in need_add_vars:
            if (
                var not in previous_block_private
                and var not in previous_block_entrance
            ):
                previous_block_entrance.append(var)
            previous_block_exit.append(var)
    return block_var_detail


def delete_block_useless_exit(
    program, program_block_ops_list, block_var_detail
):
    # forward
    for index in range(len(block_var_detail)):
        if index == len(block_var_detail) - 1:
            break
        current_block_exit = block_var_detail[index]["forward"]["exit"]
        next_block_entrance = block_var_detail[index + 1]["forward"]["entrance"]
        need_delete_var = []
        for var in current_block_exit:
            if var not in next_block_entrance:
                need_delete_var.append(var)

        for var in need_delete_var:
            current_block_exit.remove(var)
    # backward
    for index in range(len(block_var_detail) - 1, -1, -1):
        if index - 1 < 0:
            break
        current_block_exit = block_var_detail[index]["backward"]["exit"]
        next_block_entrance = block_var_detail[index - 1]["backward"][
            "entrance"
        ]
        need_delete_var = []
        for var in current_block_exit:
            if var not in next_block_entrance:
                need_delete_var.append(var)
        for var in need_delete_var:
            current_block_exit.remove(var)

    return block_var_detail


def get_communicate_var_info(
    program, block_index, entrance_var_list, type="forward"
):
    input_var_reshape_dim = []
    input_var_reshape_name = []

    if type == "forward":
        block_input_var_name = (
            f"forward_joint_{block_index - 1}_{block_index}@Heter"
        )
    else:
        block_input_var_name = (
            f"backward_joint_{block_index + 1}_{block_index}@Heter"
        )

    entrance_var_list.sort()
    # input
    # Heter_SERVER_BLOCK_index@JOINT_VAR -> slice -> var@Heter_SERVER_BLOCK@INPUT_RESHAPE_VAR -> reshape -> var
    for name in entrance_var_list:
        var = program.global_block().vars[name]
        shape = var.shape
        recv_var_dim = -1 * reduce(lambda x, y: x * y, shape, 1)
        input_var_reshape_dim.append(recv_var_dim)
        input_var_reshape_name.append(f"{name}.input_reshape@Heter")

    info = {
        "input_var_reshape_dim": input_var_reshape_dim,
        "input_var_reshape_name": input_var_reshape_name,
        "block_input_var_name": block_input_var_name,
    }

    return info


def add_vars_by_var_list(var_name_list, origin_program, program, block):
    for var_name in var_name_list:
        if (
            var_name not in program.global_block().vars
            and var_name not in block.vars
        ):
            var = origin_program.global_block().vars[var_name]
            if var.persistable:
                program.global_block()._clone_variable(
                    var, force_persistable=False
                )
            else:
                block._clone_variable(var, force_persistable=False)


def _get_output_map_from_op(varmap, op):
    """Returns a dict from op output name to the vars in varmap."""
    iomap = collections.OrderedDict()
    for key in op.output_names:
        vars = []
        for varname in op.output(key):
            if varname == "@EMPTY@":
                continue
            if "lod_tensor_blocking_queue" in varname:
                continue
            vars.append(varmap[varname])
        if len(vars) == 1:
            iomap[key] = vars[0]
        else:
            iomap[key] = vars
    return iomap


def get_varlist_from_op_map(var_map):
    var_list = []
    for key, varlist in var_map.items():
        if not isinstance(varlist, list):
            varlist = [varlist]
        for i in range(len(varlist)):
            var = varlist[i]
            var_list.append(var.name)
    return var_list


def _get_input_map_from_op(varmap, op):
    """Returns a dict from op input name to the vars in varmap."""
    iomap = collections.OrderedDict()
    for key in op.input_names:
        vars = []
        for varname in op.input(key):
            if varname == "@EMPTY@":
                continue
            if "lod_tensor_blocking_queue" in varname:
                continue
            vars.append(varmap[varname])
        if len(vars) == 1:
            iomap[key] = vars[0]
        else:
            iomap[key] = vars
    return iomap


def screen_persistables(program, var_list):
    need_remove = []
    for var_name in var_list:
        if "@GRAD" in var_name:
            if "GRAD" != var_name.split("@")[-1]:
                continue
            origin_var_name = var_name.split("@GRAD")[0]
            var = program.global_block().vars[origin_var_name]
        else:
            var = program.global_block().vars[var_name]

        if is_persistable(var):
            need_remove.append(var_name)

    for var_name in need_remove:
        var_list.remove(var_name)
    return need_remove


def block_append_op(program, origin_program, block, op):
    merge_ordereddict = origin_program.global_block().vars.copy()
    merge_ordereddict.update(block.vars)
    inputs = _get_input_map_from_op(merge_ordereddict, op)
    for key, varlist in inputs.items():
        if not isinstance(varlist, list):
            varlist = [varlist]
        for var in varlist:
            if (
                var.name not in program.global_block().vars
                and var.name not in block.vars
            ):
                if var.persistable:
                    program.global_block()._clone_variable(
                        var, force_persistable=False
                    )
                else:
                    block._clone_variable(var, force_persistable=False)

    outputs = _get_output_map_from_op(origin_program.global_block().vars, op)
    for key, varlist in outputs.items():
        if not isinstance(varlist, list):
            varlist = [varlist]
        for var in varlist:
            if (
                var.name not in program.global_block().vars
                and var.name not in block.vars
            ):
                if var.persistable:
                    program.global_block()._clone_variable(
                        var, force_persistable=False
                    )
                else:
                    block._clone_variable(var, force_persistable=False)

    if "_grad" not in op.type:
        # for forward op
        return block.append_op(
            type=op.type, inputs=inputs, outputs=outputs, attrs=op.all_attrs()
        )
    else:
        # for grad op
        op_desc = op.desc
        backward = core.op_proto_and_checker_maker.OpRole.Backward
        device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName()

        # append grad op
        new_op_desc = block.desc.append_op()
        new_op_desc.copy_from(op_desc)
        new_op_desc._set_attr(RPC_OP_ROLE_ATTR_NAME, backward)

        # set device grad
        if op.desc.has_attr(device_attr_name):
            op_device = op_desc.attr(device_attr_name)
            new_op_desc._set_attr(device_attr_name, op_device)
        block._sync_with_cpp()


def get_next_stage_trainers(role_maker):
    try:
        return role_maker._get_next_trainers()
    except Exception:
        return role_maker.get_next_trainers()


def insert_communicate_op(
    origin_program,
    role_maker,
    heter_block,
    stage_id,
    first_op_index,
    block_var_detail,
    device,
    is_forward=True,
):
    if is_forward:
        next_heter_worker_endpoints = get_next_stage_trainers(role_maker)
        previous_heter_worker_endpoints = get_previous_stage_trainers(
            role_maker
        )
        entrance_var = block_var_detail[stage_id]["forward"]["entrance"]
        comm_info = get_communicate_var_info(
            origin_program, stage_id + 1, entrance_var
        )

    else:
        next_heter_worker_endpoints = get_next_stage_trainers(role_maker)
        previous_heter_worker_endpoints = get_previous_stage_trainers(
            role_maker
        )
        entrance_var = block_var_detail[stage_id - 1]["backward"]["exit"]
        comm_info = get_communicate_var_info(
            origin_program, stage_id - 1, entrance_var, "backward"
        )

    heter_block._insert_op(
        index=first_op_index,
        type="send_and_recv",
        inputs={"X": heter_block.vars[entrance_var[0]]},
        outputs={"Out": []},
        attrs={
            "mode": "forward" if is_forward else "backward",
            "send_var_name": [*entrance_var, "microbatch_id"],
            "recv_var_name": [],
            "message_name": comm_info["block_input_var_name"],
            "next_endpoints": next_heter_worker_endpoints,
            "previous_endpoints": previous_heter_worker_endpoints,
            "trainer_id": get_role_id(role_maker),
            "op_device": device,
            RPC_OP_ROLE_ATTR_NAME: RPC_OP_ROLE_ATTR_VALUE,
        },
    )

    return entrance_var


def get_the_one_recv_context(context, is_dense=True, split_dense_table=False):
    recv_id_maps = {}
    grad_name_to_param_name = {}
    if is_dense:
        send_ctx = get_the_one_send_context(
            context, split_dense_table=split_dense_table
        )
        for idx, (name, ctx) in enumerate(send_ctx.items()):
            if ctx.is_sparse():
                continue
            if ctx.is_tensor_table():
                continue

            origin_grad_varnames = ctx.origin_varnames()

            param_names = []
            for grad_varname in origin_grad_varnames:
                param_name = context["grad_name_to_param_name"][grad_varname]
                param_names.append(param_name)
            recv_id_maps[ctx.table_id()] = param_names
    else:
        send_ctx = get_the_one_send_context(
            context, split_dense_table=False, ep_list=None
        )
        for idx, (name, ctx) in enumerate(send_ctx.items()):
            if not ctx.is_sparse():
                continue

            origin_grad_varnames = ctx.origin_varnames()

            param_names = []
            for grad_varname in origin_grad_varnames:
                param_name = context["grad_name_to_param_name"][grad_varname]
                param_names.append(param_name)
            recv_id_maps[ctx.table_id()] = param_names
    return recv_id_maps


def _get_varname_parts(varname):
    # returns origin, blockid, trainerid
    orig_var_name = ""
    trainer_part = ""
    block_part = ""
    trainer_idx = varname.find(".trainer_")
    if trainer_idx >= 0:
        trainer_part = varname[trainer_idx + 1 :]
    else:
        trainer_idx = len(varname)
    block_index = varname.find(".block")
    if block_index >= 0:
        block_part = varname[block_index + 1 : trainer_idx]
    else:
        block_index = len(varname)
    orig_var_name = varname[0 : min(block_index, trainer_idx)]
    return orig_var_name, block_part, trainer_part


dtype_to_size = {
    core.VarDesc.VarType.FP16: 2,
    core.VarDesc.VarType.FP32: 4,
    core.VarDesc.VarType.FP64: 8,
    core.VarDesc.VarType.INT16: 2,
    core.VarDesc.VarType.INT32: 4,
    core.VarDesc.VarType.INT64: 8,
    core.VarDesc.VarType.BOOL: 1,
    core.VarDesc.VarType.UINT8: 1,
}


def get_var_mem_size(var):
    m_size = reduce(lambda x, y: x * y, var.shape, 1)
    m_size *= dtype_to_size[var.dtype]
    return m_size


class MergedVariable:
    def __init__(self, merged, ordered, offsets):
        self.merged_var = merged
        self.ordered_vars = ordered
        self.offsets = offsets


def build_var_distributed(context):
    origin_programs = context['origin_main_programs']

    param_name_to_grad_name = {}
    grad_name_to_param_name = {}
    context["origin_sparse_pairs"] = []
    context["origin_dense_pairs"] = []
    context["merged_sparse_pairs"] = []
    context['merged_dense_pairs'] = []
    context["merged_variables_pairs"] = []
    context["merged_variable_map"] = {}
    for origin_program in origin_programs:
        sparse_pairs, dense_pairs = get_param_grads(origin_program)
        # print("public build_var_distributed sparse_pairs:", sparse_pairs)
        # print("public build_var_distributed dense_pairs:", dense_pairs)
        origin_for_sparse = []
        origin_for_dense = []
        merged_sparse_pairs = []
        merged_dense_pairs = []
        merged_variables_pairs = []

        for param, grad in sparse_pairs:
            origin_for_sparse.append((param, grad))

        for param, grad in dense_pairs:
            origin_for_dense.append((param, grad))

        for dense_pair in origin_for_dense:
            param, grad = dense_pair

            m_param = MergedVariable(param, [param], [0])
            m_grad = MergedVariable(grad, [grad], [0])
            merged_variables_pairs.append((m_param, m_grad))
            merged_dense_pairs.append((m_param, m_grad))
        # print("public build_var_distributed merged_dense_pairs:",
        #       merged_dense_pairs)

        for sparse_pair in origin_for_sparse:
            param, grad = sparse_pair

            m_param = MergedVariable(param, [param], [0])
            m_grad = MergedVariable(grad, [grad], [0])
            merged_variables_pairs.append((m_param, m_grad))
            merged_sparse_pairs.append((m_param, m_grad))
        # print("public build_var_distributed merged_sparse_pairs:",
        #       merged_sparse_pairs)

        for merged in merged_variables_pairs:
            m_param, m_grad = merged
            context["merged_variable_map"][m_param.merged_var.name] = (
                m_param.merged_var
            )
            context["merged_variable_map"][m_grad.merged_var.name] = (
                m_grad.merged_var
            )

        param_merges = []
        param_merges.extend(origin_for_sparse)
        param_merges.extend(origin_for_dense)

        for param, grad in param_merges:
            param_name_to_grad_name[param.name] = grad.name
            grad_name_to_param_name[grad.name] = param.name

        context["origin_sparse_pairs"].append(origin_for_sparse)
        context["origin_dense_pairs"].append(origin_for_dense)
        context["merged_sparse_pairs"].append(merged_sparse_pairs)
        context['merged_dense_pairs'].append(merged_dense_pairs)

    context["param_name_to_grad_name"] = param_name_to_grad_name
    context["grad_name_to_param_name"] = grad_name_to_param_name
    '''
    print("public build_var_distributed origin_sparse_pairs:",
        context["origin_sparse_pairs"])
    print("public build_var_distributed origin_for_dense:",
        context["origin_dense_pairs"])
    print("public build_var_distributed merged_sparse_pairs:",
        context["merged_sparse_pairs"])
    print("public build_var_distributed merged_dense_pairs:",
        context['merged_dense_pairs'])
    print("public build_var_distributed param_name_to_grad_name:",
        param_name_to_grad_name)
    print("public build_var_distributed grad_name_to_param_name:",
        grad_name_to_param_name)
    '''


def _is_opt_role_op(op):
    # NOTE : depend on oprole to find out whether this op is for
    # optimize
    op_maker = core.op_proto_and_checker_maker
    optimize_role = core.op_proto_and_checker_maker.OpRole.Optimize
    if op_maker.kOpRoleAttrName() in op.attr_names and int(
        op.all_attrs()[op_maker.kOpRoleAttrName()]
    ) == int(optimize_role):
        return True
    return False


def get_param_grads(origin_program):
    def _get_params_grads(sparse_varnames):
        block = origin_program.global_block()

        dense_param_grads = []
        sparse_param_grads = []

        optimize_params = set()
        origin_var_dict = origin_program.global_block().vars
        role_id = int(core.op_proto_and_checker_maker.OpRole.Backward)
        for op in block.ops:
            if _is_opt_role_op(op):
                # delete clip op from opt_ops when run in Parameter Server mode
                if (
                    OP_NAME_SCOPE in op.all_attrs()
                    and CLIP_OP_NAME_SCOPE in op.attr(OP_NAME_SCOPE)
                ):
                    op._set_attr("op_role", role_id)
                    continue
                if not op.has_attr(OP_ROLE_VAR_ATTR_NAME):
                    continue
                if op.attr(OP_ROLE_VAR_ATTR_NAME):
                    param_name = op.attr(OP_ROLE_VAR_ATTR_NAME)[0]
                    grad_name = op.attr(OP_ROLE_VAR_ATTR_NAME)[1]
                    if param_name not in optimize_params:
                        optimize_params.add(param_name)
                        param_grad = (
                            origin_var_dict[param_name],
                            origin_var_dict[grad_name],
                        )

                        if param_name in sparse_varnames:
                            sparse_param_grads.append(param_grad)
                        else:
                            dense_param_grads.append(param_grad)
        return sparse_param_grads, dense_param_grads

    def _get_sparse_varnames():
        varnames = []
        for op in origin_program.global_block().ops:
            if (
                op.type in SPARSE_OP_TYPE_DICT.keys()
                and op.attr('remote_prefetch') is True
            ):
                param_name = op.input(SPARSE_OP_TYPE_DICT[op.type])[0]
                varnames.append(param_name)

        return list(set(varnames))

    sparse_varnames = _get_sparse_varnames()
    sparse_param_grads, dense_param_grads = _get_params_grads(sparse_varnames)

    return sparse_param_grads, dense_param_grads


def delete_ops(block, ops):
    for op in ops:
        try:
            idx = list(block.ops).index(op)
            block._remove_op(idx)
        except Exception as e:
            print(e)


def find_send_op(program):
    send_op_list = []
    for op in program.global_block().ops:
        if op.type == "send":
            send_op_list.append(op)
    return send_op_list


def find_op_input_output(program, block, op):
    input_var_list = []
    output_var_list = []
    inputs = _get_input_map_from_op(block.vars, op)
    input_var_list += get_varlist_from_op_map(inputs)
    outputs = _get_output_map_from_op(block.vars, op)
    output_var_list += get_varlist_from_op_map(outputs)
    input_var_list = list(set(input_var_list))
    output_var_list = list(set(output_var_list))
    return input_var_list, output_var_list


def add_send_op(program, block, _vars):
    def _get_send_op_dict():
        send_op_dict = {}
        send_op_list = find_send_op(program)
        for op in send_op_list:
            input_list, _ = find_op_input_output(
                program, program.global_block(), op
            )
            for var in input_list:
                send_op_dict[var] = op
        return send_op_dict

    send_grad_var_list = []
    send_op_dict = _get_send_op_dict()
    table_dict = {}
    for persistable_var in _vars:
        if "@GRAD" not in persistable_var:
            continue
        if "GRAD" != persistable_var.split("@")[-1]:
            continue
        if persistable_var not in send_op_dict:
            continue
        send_op = send_op_dict[persistable_var]
        is_sparse = send_op.attr('is_sparse')
        table_id = send_op.attr('table_id')
        send_varnames = send_op.attr('send_varnames')
        send_grad_var_list.append(persistable_var)
        if table_id not in table_dict:
            table_dict[table_id] = {}
            table_dict[table_id]['var_list'] = []
            table_dict[table_id]['is_sparse'] = is_sparse
            table_dict[table_id]['send_varnames'] = send_varnames
        table_dict[table_id]['var_list'].append(persistable_var)

    for table_id in table_dict:
        dummy_output = block.create_var(name=generate_control_dev_var_name())
        send_input_vars = [
            block.vars[union_var]
            for union_var in table_dict[table_id]['var_list']
        ]
        block.append_op(
            type="send",
            inputs={"X": send_input_vars},
            outputs={"Out": dummy_output},
            attrs={
                "send_varnames": table_dict[table_id]['send_varnames'],
                "is_sparse": is_sparse,
                "table_id": table_id,
                RPC_OP_ROLE_ATTR_NAME: RPC_OP_ROLE_ATTR_VALUE,
            },
        )

    return send_grad_var_list


def get_vars_name_in_block(block):
    vars_list = block.vars.keys()
    vars_name_list = list(vars_list)
    return vars_name_list


# reserve static_var
def delete_trainer_useless_var(program, static_var):
    static_var = list(set(static_var))
    program_useful_var_list = []
    for op in program.global_block().ops:
        input_var_list, output_var_list = find_op_input_output(
            program, program.global_block(), op
        )
        op_var_list = list(set(input_var_list).union(set(output_var_list)))
        program_useful_var_list = list(
            set(program_useful_var_list).union(set(op_var_list))
        )
    program_useful_var_list += static_var
    program_useless_var_list = list(
        set(get_vars_name_in_block(program.global_block())).difference(
            set(program_useful_var_list)
        )
    )
    for var in program_useless_var_list:
        program.global_block()._remove_var(var)
    return program_useless_var_list


def create_backward_block(
    program, origin_program, bp_ops_list, block_var_detail
):
    pre_block_idx = program.num_blocks - 1
    heter_block = program._create_block(pre_block_idx)

    for _, op in enumerate(bp_ops_list):
        if op.type == "send":
            send_varnames = op.attr('send_varnames')
            is_skip = False
            for varname in send_varnames:
                if (
                    varname not in program.global_block().vars
                    and varname not in heter_block.vars
                ):
                    is_skip = True
                    break
            if is_skip:
                continue
        block_append_op(program, origin_program, heter_block, op)

    entrance_vars = block_var_detail[0]["backward"]["entrance"]
    add_vars_by_var_list(entrance_vars, origin_program, program, heter_block)
    exit_vars = block_var_detail[0]["backward"]["exit"]
    add_vars_by_var_list(exit_vars, origin_program, program, heter_block)
    return heter_block


def is_backward_op(op):
    return op_role_attr_name in op.attr_names and (
        int(op.attr(op_role_attr_name)) & int(op_role.Backward)
    )


def is_forward_op(op):
    return op_role_attr_name in op.attr_names and (
        int(op.attr(op_role_attr_name)) == int(op_role.Forward)
    )


def is_push_sparse_op(op):
    return op.type == 'distributed_push_sparse'


def get_distributed_push_sparse_op_list(block):
    push_sparse_op_list = []
    for op_idx in range(block.desc.op_size()):
        op = block.ops[op_idx]
        if is_push_sparse_op(op):
            push_sparse_op_list.append(op)
    return push_sparse_op_list


def get_bp_op_list(block):
    bp_op_list = []
    for op_idx in range(block.desc.op_size()):
        op = block.ops[op_idx]
        if is_backward_op(op):
            bp_op_list.append(op)
    return bp_op_list


def delete_same_ops(block, ops):
    for op in ops:
        try:
            for origin_op in block.ops:
                if str(origin_op) == str(op):
                    idx = list(block.ops).index(origin_op)
                    block._remove_op(idx)
                    break
        except Exception as e:
            print(e)


def check_program(program):
    block_idx = 0
    for block in program.blocks:
        for op in block.ops:
            input_var_names = op.desc.input_arg_names()
            output_var_names = op.desc.output_arg_names()
            for var_name in input_var_names + output_var_names:
                if not block._find_var_recursive(str(var_name)):
                    raise ValueError(
                        f'var: {var_name} needed by op is not found in block: {block_idx}'
                    )
        block_idx += 1
    print('program checked valid')


def debug_program(file, program):
    # py >= 3.2
    os.makedirs(os.path.dirname(file), exist_ok=True)
    with open(file, 'w+') as f:
        f.write(str(program))


def is_distributed_env():
    node_role = os.getenv("TRAINING_ROLE")
    if node_role is None:
        return False
    else:
        return True