import json
import os

from corekit.hikframe.dl.utils.env import print_envs


class OptRecorder(object):
    @classmethod
    def __init__(
        cls, save_dir=None, status_list=None,
        step_status_list=None, step_names=None
        ):
        # 显示传入的env参数
        cls.envs = ['save_dir']
        print_envs(cls.envs)

        if not step_names:
            cls.step_names = step_names
        else:
            cls.step_names = ['training', 'adapting']

        # 获取保存路径
        if save_dir is not None:
            cls.save_dir = save_dir
        else:
            cls.save_dir = os.getenv("save_dir")
        if cls.save_dir is None:
            raise Exception("未指定保存目录. 请设置环境变量'save_dir'或者输入初始化参数 'save_dir'.")

        # save_model_dir = os.path.join(save_dir, 'model')
        save_config_dir = os.path.join(cls.save_dir, 'config')
        os.makedirs(save_config_dir, exist_ok=True)
        cls.json_path = os.path.join(save_config_dir, 'opt_rec.json')

        if not status_list:
            cls.status_list = ['RUNNING', 'COMPLETED']
        if not step_status_list:
            cls.step_status_list = ['RUNNING', 'COMPLETED']
        cls.trial_init_dict = {
            "status": "",
            "flow_step": {},
            "index": {}
        }
        cls._init_record_info()

    @classmethod
    def _init_record_info(cls):
        os.makedirs(os.path.dirname(cls.json_path), exist_ok=True)
        if not os.path.isfile(cls.json_path):
            cls.record_info = {
                "best_trial": "-1",
                "current_trial": "-1",
                "-1": cls.trial_init_dict
            }
            cls._write(cls.record_info)
        else:
            cls.record_info = cls._read()
        print('[OPERATION] init_record_info')
        print('=' * 10, 'record_info')
        print(cls.record_info)

    @classmethod
    def _read(cls):
        with open(cls.json_path, 'r') as f:
            info = json.load(f)
            return info

    @classmethod
    def _write(cls, info):
        with open(cls.json_path, 'w') as f:
            json.dump(info, f)

    @classmethod
    def start_trial(cls):
        info = cls._read()
        current_trial = str(int(info.get("current_trial")) + 1)
        info["current_trial"] = current_trial
        info[current_trial] = cls.trial_init_dict
        info[current_trial]["status"] = cls.status_list[0]
        print('[OPERATION] start_trial, current_trial_id is [{}]'.format(
            current_trial))
        print('info[current_trial]')
        print(info[current_trial])
        cls._write(info)
        return current_trial

    @classmethod
    def stop_trial(cls, trial_id):
        info = cls._read()
        current_trial = trial_id
        info[current_trial]["status"] = cls.status_list[-1]
        print('[OPERATION] stop_trial, current_trial_id is [{}]'.format(
            current_trial))
        print('info[current_trial]')
        print(info[current_trial])
        cls._write(info)
        return

    @classmethod
    def start_step(cls, trial_id, step_name):
        info = cls._read()
        current_trial = trial_id
        info[current_trial]["flow_step"].update(
            {step_name: cls.step_status_list[0]})
        print('[OPERATION] start_step, current_trial_id is [{}], step_name is [{}]'.format(
            current_trial, step_name))
        print('info[current_trial]')
        print(info[current_trial])
        cls._write(info)
        return

    @classmethod
    def stop_step(cls, trial_id, step_name):
        info = cls._read()
        current_trial = trial_id
        info[current_trial]["flow_step"].update(
            {step_name: cls.step_status_list[-1]})
        print('[OPERATION] stop_step, current_trial_id is [{}], step_name is [{}]'.format(
            current_trial, step_name))
        print('info[current_trial]')
        print(info[current_trial])
        cls._write(info)
        return

    @classmethod
    def update_metric(cls, trial_id, metric):
        info = cls._read()
        current_trial = trial_id
        info[current_trial]["index"].update(metric)
        print('[OPERATION] update_metric, current_trial_id is [{}], metric is [{}]'.format(
            current_trial, metric))
        best_trial_index = info[info["best_trial"]]["index"]
        if not best_trial_index or (float(list(metric.values())[0]) > float(list(best_trial_index.values())[0])):
            print('best trial updated. from {} to {}'.format(
                best_trial_index, metric))
            info["best_trial"] = trial_id
        print('info[current_trial]')
        print(info[current_trial])
        cls._write(info)
        return

# if __name__ == '__main__':
#     opt_rec_json = './xx/test.json'
#     TRAIN_STEP = 'training'
#     ADAPTE_STEP = 'adapting'

#     for i in range(3):
#         recorder = OptRecorder(opt_rec_json)
#         cur_trial_id = recorder.start_trial()
#         recorder.start_step(trial_id=cur_trial_id, step_name=TRAIN_STEP)
#         recorder.stop_step(trial_id=cur_trial_id, step_name=TRAIN_STEP)
#         # adapt
#         recorder.start_step(trial_id=cur_trial_id, step_name=ADAPTE_STEP)
#         recorder.stop_step(trial_id=cur_trial_id, step_name=ADAPTE_STEP)

#         # todo，更新metrics
#         train_metric = {'map': 0.28+i}
#         recorder.update_metric(trial_id=cur_trial_id, metric=train_metric)

#         recorder.stop_trial(trial_id=cur_trial_id)
