# Copyright (c) 2020 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.

"""
Training use fluid with one node only.
"""

from __future__ import print_function

import logging

import numpy as np
import paddle.fluid as fluid
from paddle.fluid.incubate.fleet.parameter_server.distribute_transpiler import fleet

from paddlerec.core.utils import envs
from paddlerec.core.trainers.cluster_trainer import ClusterTrainer

logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger("fluid")
logger.setLevel(logging.INFO)
special_param = ["TDM_Tree_Travel", "TDM_Tree_Layer", "TDM_Tree_Info"]


class TDMClusterTrainer(ClusterTrainer):
    def server(self, context):
        namespace = "train.startup"
        init_model_path = envs.get_global_env(
            "cluster.init_model_path", "", namespace)
        assert init_model_path != "", "Cluster train must has init_model for TDM"
        fleet.init_server(init_model_path)
        logger.info("TDM: load model from {}".format(init_model_path))
        fleet.run_server()
        context['is_exit'] = True

    def startup(self, context):
        self._exe.run(fleet.startup_program)

        namespace = "train.startup"
        load_tree = envs.get_global_env(
            "tree.load_tree", True, namespace)
        self.tree_layer_path = envs.get_global_env(
            "tree.tree_layer_path", "", namespace)
        self.tree_travel_path = envs.get_global_env(
            "tree.tree_travel_path", "", namespace)
        self.tree_info_path = envs.get_global_env(
            "tree.tree_info_path", "", namespace)

        save_init_model = envs.get_global_env(
            "cluster.save_init_model", False, namespace)
        init_model_path = envs.get_global_env(
            "cluster.init_model_path", "", namespace)

        if load_tree:
            # covert tree to tensor, set it into Fluid's variable.
            for param_name in special_param:
                param_t = fluid.global_scope().find_var(param_name).get_tensor()
                param_array = self._tdm_prepare(param_name)
                param_t.set(param_array.astype('int32'), self._place)

        if save_init_model:
            logger.info("Begin Save Init model.")
            fluid.io.save_persistables(
                executor=self._exe, dirname=init_model_path)
            logger.info("End Save Init model.")

        context['status'] = 'train_pass'

    def _tdm_prepare(self, param_name):
        if param_name == "TDM_Tree_Travel":
            travel_array = self._tdm_travel_prepare()
            return travel_array
        elif param_name == "TDM_Tree_Layer":
            layer_array, _ = self._tdm_layer_prepare()
            return layer_array
        elif param_name == "TDM_Tree_Info":
            info_array = self._tdm_info_prepare()
            return info_array
        else:
            raise " {} is not a special tdm param name".format(param_name)

    def _tdm_travel_prepare(self):
        """load tdm tree param from npy/list file"""
        travel_array = np.load(self.tree_travel_path)
        logger.info("TDM Tree leaf node nums: {}".format(
            travel_array.shape[0]))
        return travel_array

    def _tdm_layer_prepare(self):
        """load tdm tree param from npy/list file"""
        layer_list = []
        layer_list_flat = []
        with open(self.tree_layer_path, 'r') as fin:
            for line in fin.readlines():
                l = []
                layer = (line.split('\n'))[0].split(',')
                for node in layer:
                    if node:
                        layer_list_flat.append(node)
                        l.append(node)
                layer_list.append(l)
        layer_array = np.array(layer_list_flat)
        layer_array = layer_array.reshape([-1, 1])
        logger.info("TDM Tree max layer: {}".format(len(layer_list)))
        logger.info("TDM Tree layer_node_num_list: {}".format(
            [len(i) for i in layer_list]))
        return layer_array, layer_list

    def _tdm_info_prepare(self):
        """load tdm tree param from list file"""
        info_array = np.load(self.tree_info_path)
        return info_array