import ray
from ray import tune
from ray.rllib.models import ModelCatalog
from ray.rllib.utils.framework import try_import_tf
from ray.tune import registry
import shutil
import os
import json
from corekit.hikframe.rl.base import RLAlgBase
from .core.policy import init_policy
from .core.trainer import init_trainer, get_default_config, set_weights, get_weights
from .core.config import config_from_file, Config
from ray.tune.logger import pretty_print
from .core.utils.log import summary_write, summary_close
import corekit.hikframe


class HikRL(RLAlgBase):
    @classmethod
    def init(cls, num_cpus=0, num_gpus=0):
        print('HIKFRAME VERSION:', corekit.hikframe.__version__)
        super().init()
        tf1, tf, tfv = try_import_tf()
        ray.init(num_cpus=num_cpus, num_gpus=num_gpus)
        cls.num_cpus = num_cpus
        cls.num_gpus = num_gpus
        assert ray.is_initialized() == True

    @classmethod
    def shutdown(cls):
        ray.shutdown()
        assert ray.is_initialized() == False

    @classmethod
    def download_dataset(cls, path):
        return super().download_dataset(path)

    @classmethod
    def get_dataset_dir(cls):
        return super().get_dataset_dir()

    @classmethod
    def get_dev_dir(cls):
        if hasattr(cls, 'save_dir'):
            cls.tmp_save_dir = os.path.join(
                cls.save_dir, "train_results")
        else:
            print('[ERROR] can not get save_dir from env.')
            cls.tmp_save_dir = ''
        try:
            if os.path.exists(cls.tmp_save_dir):
                shutil.rmtree(cls.tmp_save_dir)
        except Exception as e:
            print(e)
        os.makedirs(cls.tmp_save_dir, exist_ok=True)
        return cls.tmp_save_dir

    @classmethod
    def get_log_dir(cls):
        if hasattr(cls, 'save_log_dir'):
            return cls.save_log_dir
        else:
            print('[ERROR] can not get save_log_dir from env.')
            cls.save_log_dir = ''

    @classmethod
    def get_pretrained_trainer(cls):
        restore_model_path = os.getenv('restore_model_path')
        ckpt_path = None
        if restore_model_path:
            pretrained_path = os.path.join(
                os.getenv('restore_model_path'), 'model')
            if os.path.exists(pretrained_path):
                ckpt_list = os.listdir(pretrained_path)
                ckpt_list.sort(key=lambda x: int(x.split('_')[1]))
                ckpt_dir_nm = ckpt_list[-1]
                ckpt_path = os.path.join(pretrained_path, ckpt_dir_nm,
                                         ckpt_dir_nm.replace('_', '-'))
                return ckpt_path
        print('[!!!] no pretrained trainer got.')
        return ckpt_path

    @classmethod
    def restore_trainer(cls, trainer, ckpt_path=None):
        if not ckpt_path:
            ckpt_path = HikRL.get_pretrained_trainer()
        if not ckpt_path or not os.path.isfile(ckpt_path):
            print('[ERROR] trainer restore failed, ckpt_path must direct to a ckpt file!, but input is {}'.format(
                ckpt_path))
            return trainer
        trainer.restore(ckpt_path)
        print('[SUCCESS] trainer restored success.')
        return trainer

    @classmethod
    def run_trainer(cls, trainer):
        result = trainer.train()
        return result

    @classmethod
    def save_trainer(cls, trainer, ckpt_dir=None):
        ckpt_dir = ckpt_dir if ckpt_dir else 'checkpoint'
        if ckpt_dir:
            os.makedirs(ckpt_dir, exist_ok=True)
        result = trainer.save(ckpt_dir)
        print('checkpoint is saved in', result)

    @classmethod
    def stop_trainer(cls, trainer):
        trainer.stop()

    @classmethod
    def register_custom_model(cls, name, mod):
        ModelCatalog.register_custom_model(name, mod)

    @classmethod
    def pretty_print(cls, result):
        return pretty_print(result)

    @classmethod
    def register_env(cls, name, env_creator):
        registry.register_env(name, env_creator)

    @classmethod
    def get_default_config(cls, name):
        return get_default_config(name)

    @classmethod
    def config_from_file(cls, config_file):
        """
        print(cfg._cfg_dict)
        # 格式化打印
        print(cfg.pretty_text)
        """
        return config_from_file(config_file)

    @classmethod
    def init_trainer(cls, name, config):
        """
        trainer = PPOTrainer(config=config, env=YourEnv)
        result= trainer.train()
        result["episode_reward_mean"]
        result["episode_reward_mean"]

        """
        return init_trainer(name, config)

    @classmethod
    def init_policy(cls, name, frame):
        return init_policy(name, frame)

    @classmethod
    def set_weights(cls, trainer, weights):
        set_weights(trainer, weights)

    @classmethod
    def get_weights(cls, trainer, policy):
        return get_weights(trainer, policy)

    @classmethod
    def report(cls, result_dict, step_num, save_log_dir):
        assert type(result_dict) == dict, 'input must be a dict'
        # cls.summary.write(result_dict, step_num)
        summary_write(result_dict, step_num, save_log_dir)
        # tune.report(**result_dict)

    @classmethod
    def train(cls, training_function, config, save_dir=None, stop=None):
        local_dir = save_dir
        os.makedirs(local_dir, exist_ok=True)
        stop = stop if stop else None
        cfg = Config(config)
        print('config info is:')
        try:
            print(cfg.pretty_text)
        except:
            print(cfg._cfg_dict)
        print('\n')

        trial_nm = 'train_exp'
        cls.analysis = ray.tune.run(
            training_function,
            # "PPO",
            config=config,
            # stop,             $ stop condition dict
            # 0, 1, 2, or 3. Verbosity mode. 0 = silent, 1 = only status updates, 2 = status and brief trial results, 3 = status and detailed trial results. Defaults to 3.
            verbose=1,
            # How many training iterations between checkpoints. A value of 0 (default) disables checkpointing. This has no effect when using the Functional Training API.
            # checkpoint_freq=1,
            # keep_checkpoints_num=3,
            # # Number of checkpoints to keep. A value of None keeps all checkpoints. Defaults to None. If set, need to provide checkpoint_score_attr.
            # checkpoint_at_end=True,
            # Specifies by which attribute to rank the best checkpoint. like "min-xxx", ""  ???还有哪些
            # checkpoint_score_attr="mean_accuracy",
            # name="test_experiment_tf",                        # Name of experiment
            local_dir=local_dir,  # Local dir to save training results to
            # restore=CURPATH+"/hikrl_results/test_experiment_tf/PPO_KZ_bf10a_00000_0_2021-03-22_22-51-00/checkpoint_2/checkpoint-2",                            # Path to checkpoint. Only makes sense to set if running 1 trial.
            # resume=True,                                      # Defaults is False. One of “LOCAL”, “REMOTE”, “PROMPT”, “ERRORED_ONLY”, or bool. LOCAL/True restores the checkpoint from the local_checkpoint_dir, determined by name and local_dir. REMOTE restores the checkpoint from remote_checkpoint_dir. PROMPT provides CLI feedback. False forces a new experiment. ERRORED_ONLY resets and reruns ERRORED trials upon resume - previous trial artifacts will be left untouched. If resume is set but checkpoint does not exist, ValueError will be thrown.
            name=trial_nm
        )
        # cls._analysis()
        cls._cosntruct_model_pkg(save_dir)
        summary_close()

    # @classmethod
    # def _analysis(cls):
    #     # best_trial = cls.analysis.best_trial  # Get best trial
    #     # best_config = cls.analysis.best_config  # Get best trial's hyperparameters
    #     # best_logdir = cls.analysis.best_logdir  # Get best trial's logdir
    #     # Get best trial's best checkpoint
    #     best_checkpoint = cls.analysis.best_checkpoint
    #     best_result = cls.analysis.best_result  # Get best trial's last results
    #     # Get best result as pandas dataframe
    #     best_result_df = cls.analysis.best_result_df
    #     # print('best_checkpoint', best_checkpoint)
    #     print('===============> ', 'best_result:')
    #     print(best_result)
    #     print('===============> ', 'best_result_df:')
    #     print(best_result_df)

    # def check(stop_reward):
    #     check_learning_achieved(cls.analysis, stop_reward)
    # @classmethod
    # def __del__(cls):
    #     ray.shutdown()
    #     assert ray.is_initialized() == False

    # @classmethod
    # def save_trainer(cls, new_dir):
    #     if os.path.exists(new_dir):
    #         shutil.rmtree(new_dir)
    #     save_dir = cls.tmp_save_dir
    #     assert save_dir and os.path.exists(save_dir)
    #     for root, dirs, files in os.walk(save_dir):
    #         for dir_name in dirs:
    #             # copy model
    #             if 'checkpoint_' in dir_name:
    #                 shutil.copytree(os.path.join(root, dir_name), new_dir)

    @classmethod
    def _cosntruct_model_pkg(cls, save_dir):
        assert save_dir and os.path.exists(save_dir)
        for root, dirs, files in os.walk(save_dir):
            for dir_name in dirs:
                # copy model
                if 'checkpoint_' in dir_name:
                    shutil.copytree(os.path.join(root, dir_name),
                                    os.path.join(cls.save_model_dir, dir_name))
            for file_name in files:
                #     # copy tensorboard file
                #     if file_name.startswith('events.out'):
                #         shutil.copy(os.path.join(root, file_name),
                #                     os.path.join(cls.save_log_dir, file_name))
                # copy model for deduce
                if file_name.endswith(('.ckpt', '.h5')):
                    shutil.copy(os.path.join(root, file_name),
                                os.path.join(cls.save_dir, file_name))
                # copy config json
                if file_name == 'params.json':
                    shutil.copy(os.path.join(root, file_name),
                                os.path.join(cls.save_cfg_dir, file_name))
                    # try:
                    #     with open(os.path.join(root, file_name)) as f:
                    #         param_json = json.load(f)
                    #         cls.config["version_info"]["hyperparameter"] = param_json
                    #         with open(os.path.join(cls.save_cfg_dir, 'modelmgt.json'), 'w+') as config_file:
                    #             json.dump(cls.config, config_file)
                    # except Exception as e:
                    #     print(e)
        # try:
        #     shutil.rmtree(cls.tmp_save_dir)
        # except Exception as e:
        #     print(e)
        print('model package constructed success.')
