import json
import os

# pip install bayesian-optimization
import numpy as np
import ray
from ray import tune
from ray.tune import Stopper
from ray.tune.suggest.basic_variant import BasicVariantGenerator
# from ray.tune.schedulers import AsyncHyperBandScheduler
from ray.tune.suggest.bayesopt import BayesOptSearch

from corekit.hikframe import HikModel
from corekit.hikframe.dl.utils.env import print_envs
from .log import summary_close, summary_write


class CustomStopper(Stopper):
    def __init__(self, target_metric, target_max_value):
        self.should_stop = False
        self.target_metric = target_metric
        self.target_max_value = target_max_value

    def __call__(self, trial_id, result):
        print('result')
        print(result)
        if not self.should_stop and result[self.target_metric] >= self.target_max_value:
            self.should_stop = True
        return self.should_stop

    def stop_all(self):
        """Returns whether to stop trials and prevent new ones from starting."""
        return self.should_stop


class HikHPS(HikModel):
    @classmethod
    def init(cls, save_dir=None):
        super().__init__(save_dir)
        # 显示传入的env参数
        cls.envs = ['search_alg', 'search_max_num',
                    'search_metric', 'search_metric_value']
        print_envs(cls.envs, 'HikHPS')

        try:
            cls.search_alg = os.environ['search_alg']
            cls.search_max_num = int(os.environ['search_max_num'])
            cls.search_metric = os.environ['search_metric']
            cls.search_metric_value = float(os.environ['search_metric_value'])
        except Exception as e:
            print(e)
            print('[ERROR] please make sure hps env is set.')
            exit()

    @classmethod
    def _parser_search_config(cls, search_space, search_config):
        metric = search_config["target_metric"]
        search_alg = search_config["search_alg"]

        try:
            alg_list = ['grid', 'random', 'bayes']
            assert search_alg in alg_list
        except:
            print('[ERROR] wrong search_alg "{}", only support "{}"'.format(
                search_alg, alg_list))
            return

        # new search_space
        new_space = {}
        for k, v in search_space.items():
            # assert 1 <= len(
            #     v) <= 2, 'search_value must be a list and num<=2'
            if isinstance(v, list):
                if len(v) == 1:
                    new_space[k] = v[0]
                elif len(v) == 2 and isinstance(v[0], float) and isinstance(v[1], float):
                    if search_alg == 'grid':
                        interval = int((v[1] - v[0]) / v[0])
                        new_space[k] = tune.grid_search(
                            np.linspace(v[0], v[1], interval).tolist())
                    elif search_alg == 'random':
                        new_space[k] = tune.uniform(v[0], v[1])
                    elif search_alg == 'bayes':
                        new_space[k] = tune.uniform(v[0], v[1])
                    else:
                        pass
                else:
                    new_space[k] = tune.choice(v)
                    pass
            else:
                new_space[k] = v
        # grid,random,bayes
        if search_alg == 'grid':
            search_alg = BasicVariantGenerator()
        elif search_alg == 'random':
            search_alg = BasicVariantGenerator()
        elif search_alg == 'bayes':
            search_alg = BayesOptSearch(metric=metric, mode="max")
        else:
            pass
        print('search_space', new_space)
        print('search_alg', search_alg)
        return new_space, search_alg

    @classmethod
    def hps_search(cls, train_func, train_config, search_config, save_dir):
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        # metric = search_config["target_metric"]
        resources_per_trial = search_config["resources_per_trial"]
        parallel_trial_num = search_config["search_max_num"]

        # stop config
        # stop = {
        #     # "training_iteration": int(search_config["search_max_num"]),
        #     search_config["target_metric"]: float(
        #         search_config["target_max_value"])
        # }
        stopper = CustomStopper(search_config["target_metric"], float(
            search_config["target_max_value"]))

        search_space, search_alg = cls._parser_search_config(
            train_config['search_space'], search_config)
        train_config['search_space'] = search_space
        local_dir = save_dir
        os.makedirs(local_dir, exist_ok=True)

        analysis = tune.run(
            train_func,
            mode="max",
            name="hps_exp",
            #         scheduler=sched,
            search_alg=search_alg,
            stop=stopper,
            resources_per_trial=resources_per_trial,
            # 并行搜索数
            num_samples=parallel_trial_num,
            config=train_config,
            local_dir=local_dir,
            verbose=2,
            # fail_fast=True,
        )

        summary_close()

        # move best trail's model and config
        save_cfg_dir = os.path.join(save_dir, 'config')
        metric_json = os.path.join(save_cfg_dir, 'hps_search.json')
        best_metric = {}
        if os.path.isfile(metric_json):
            with open(metric_json, 'r') as f:
                info = json.load(f)
            best_trail = info.get('best_trail')
            if int(best_trail) > -1:
                best_metric = info.get('search_history').get(
                    best_trail).get('eval_metric')
                trail_dir = cls.get_trail_dir(save_dir, best_trail)
                for dirname in ['model', 'config']:
                    print('cp -r {}/* {}/'.format(os.path.join(trail_dir,
                                                               dirname), os.path.join(save_dir, dirname)))
                    os.system('cp -r {}/* {}/'.format(os.path.join(trail_dir,
                                                                   dirname), os.path.join(save_dir, dirname)))

            # 删除所有trail模型避免占用过多空间
            trails = list(info['search_history'].keys())
            for trail in trails:
                if trail == best_trail:
                    continue
                trail_dir = cls.get_trail_dir(save_dir, trail)
                print('trail_dir is {}'.format(trail_dir))
                for model_subfix in ['.pth', '.pb']:
                    os.system(
                        'rm -rf {}/model/*{}'.format(trail_dir, model_subfix))
        else:
            print('{} is not exists.'.format(metric_json))
        return best_metric

    @classmethod
    def record(cls, step_num, search_space, record_metric, save_dir, model_grah_path=None):
        tb_dir = os.path.join(save_dir, 'log')
        save_cfg_dir = os.path.join(save_dir, 'config')
        if not os.path.exists(tb_dir):
            os.makedirs(tb_dir)
        if not os.path.exists(save_cfg_dir):
            os.makedirs(save_cfg_dir)

        # record tensorboard
        summary_write(record_metric, step_num, tb_dir)
        summary_write(search_space, step_num, tb_dir)

        # record hps_search history
        metric_json = os.path.join(save_cfg_dir, 'hps_search.json')
        if os.path.isfile(metric_json):
            with open(metric_json, 'r') as f:
                info = json.load(f)
        else:
            info = {'search_history': {}, 'best_trail': -1}

        info['search_history'][step_num] = {
            'eval_metric': record_metric, 'search_space': search_space, 'model_graph': model_grah_path
        }
        best_trail = info['best_trail']
        if best_trail == -1:
            best_trail = step_num
        else:
            eval_metric = info['search_history'][best_trail]['eval_metric']
            if list(record_metric.values())[0] >= list(eval_metric.values())[0]:
                best_trail = step_num
        info['best_trail'] = str(best_trail)
        print('current best trail info is: {}'.format(
            info['search_history'][best_trail]))
        with open(metric_json, 'w') as f:
            json.dump(info, f)
        print('metric json is updated.')

        tune.report(**record_metric)

    @classmethod
    def get_trail_dir(cls, save_dir, trail_id):
        trail_save_dir = os.path.join(save_dir, 'all_trail')
        os.makedirs(trail_save_dir, exist_ok=True)
        cur_trail_dir = os.path.join(
            trail_save_dir, 'trail_{}'.format(trail_id))
        os.makedirs(cur_trail_dir, exist_ok=True)
        return cur_trail_dir

    # @classmethod
    # def get_trail_id(cls, save_dir):
    #     import fcntl
    #     import threading
    #     record_txt = os.path.join(save_dir, 'trail_count.txt')
    #     last_num = -1
    #     if os.path.isfile(record_txt):
    #         with open(record_txt, 'r') as f:
    #             last_num = int(f.readlines()[0])
    #     last_num += 1
    #     with open(record_txt, 'w') as f2:
    #         f2.writelines('{}'.format(last_num))
    #     return last_num

    @classmethod
    def get_share_obj(cls, name="shared_val"):
        actor = ray.get_actor(name)
        return ray.get(actor.get.remote())

    @classmethod
    def set_share_obj(cls, name="shared_val", val=0):
        actor = ray.get_actor(name)
        actor.set.remote(val)

    @classmethod
    def init_share_obj(cls, name="shared_val", val=0):
        actor = ShareClass.options(name=name).remote()
        actor.set.remote(val)
        return actor

    @classmethod
    def obj_share(cls, actor):
        pass


@ray.remote
class ShareClass:
    def __init__(self):
        self.value = 0

    def get(self):
        return self.value

    def set(self, value):
        self.value = value
