import json
import os
import shutil
import time

import numpy as np

# from corekit.hikframe.dl.utils.env import print_envs
from ..utils.filesys import copyfiles, search_files

try:
    import tensorboardX as tbX
except ImportError:
    print('WARNING: no tensorboardX installed, tensorboard file may not be saved.')
    tbX = None


class HikModel:
    """
    训练模型集成工具
    """
    envs = ['gpu_num', 'data_dir', 'code_dir', 'save_dir',
            'pretrain', 'resume', 'restore_model_path', 'encrypt', 'encrypt_key']
    save_dir = os.getenv("save_dir", '')
    '''模型包包含`log`, `model`，`val`和`config`四部分.
    
    log 中包含用于可视化模型结构及结果的tensorboard 文件
    model 包含训练输出的模型文件
    val 包含模型评估对每个输出数据的预测结果
    config 包含该模型版本的基础信息和在其创建的模型评估、模型压缩、模型转换任务
    '''
    save_model_dir = os.path.join(save_dir, "model")
    save_log_dir = os.path.join(save_dir, "log")
    save_cfg_dir = os.path.join(save_dir, "config")
    save_val_dir = os.path.join(save_dir, "val")

    # tensorbordX 写入
    writer = tbX.SummaryWriter(save_log_dir) if tbX else None
    config = {}
    encrypt_key = None
    config_path: str
    dataset_dir: str
    agents_dir: str  # 写入联赛训练信息的路径

    @classmethod
    def __init__(cls, save_dir=None, data_dir=None):
        """
        初始化函数
        """
        # 显示传入的env参数
        cls.envs = ['gpu_num', 'data_dir', 'code_dir', 'save_dir',
                    'pretrain', 'resume', 'restore_model_path', 'encrypt', 'encrypt_key']
        # print_envs(cls.envs)
        # 初始化模型目录
        cls._init_model_package(save_dir)

        # tensorbordX 写入
        cls.writer = tbX.SummaryWriter(cls.save_log_dir) if tbX else None

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

        # 获取任务参数
        cls.code_dir = os.getenv("code_dir", "")
        cls.gpu_num = int(os.getenv("gpu_num", 0))
        cls.task_name = os.getenv("task_name", "")
        cls.pretrain_path = os.getenv("restore_model_path", "")
        cls.encrypt = os.getenv("encrypt", "")
        cls.encrypt_key = os.getenv("encrypt_key", "")
        # 初始化模型配置
        cls.config = {
            "user": "",
            "preset": "",
            "name": "",
            "usage": "",
            "engine": "",
            "scene": "",
            "model_type": "",
            "description": "",
            "version_info":
                {
                    "version": "",
                    "tag": "",
                    "model_path": "",
                    "algorithm_name": "",
                    "algorithm_path": "",
                    "framework": "",
                    "image": {},
                    "dataset_type": "",
                    "dataset_usage": "",
                    "dataset_name": "",
                    "dataset_path": "",
                    "create_time": "",
                    "resource_type": "",
                    "resource_num": "",
                    "hyperparameter": {},
                    "log_path": "",
                    "model_size": "",
                    "isdeploy": "",
                    "train_metrics": {},
                    "val_metrics": {},
                    "encrypt_key": cls.encrypt_key
                }
        }

    @classmethod
    def _init_model_package(cls, save_dir):
        """
        初始化模型包的格式

        模型包包含`log`, `model`，`val`和`config`四部分，其中,
        `log`中包含用于可视化模型结构及结果的`tensorboard`文件，
        `model`包含训练输出的模型文件，
        `val`包含模型评估对每个输出数据的预测结果，
        `config`包含该模型版本的基础信息和在其创建的模型评估、模型压缩、模型转换任务。
        """
        # 获取保存路径
        # if 'win' in sys.platform:
        #     cls.save_dir = save_dir  # 本地下采用输入参数
        # else:
        #     cls.save_dir = os.getenv("save_dir")  # 服务器上采用环境变量
        cls.save_dir = os.getenv("save_dir", save_dir)  # 2024.09.14 为测试linux下的分布式训练, 临时改动

        if cls.save_dir is None:
            raise Exception("未指定保存目录. 请设置环境变量'save_dir'或者输入初始化参数 'save_dir'.")

        # 创建模型目录
        cls.save_model_dir = os.path.join(cls.save_dir, "model")
        cls.save_log_dir = os.path.join(cls.save_dir, "log")
        cls.save_cfg_dir = os.path.join(cls.save_dir, "config")
        cls.save_val_dir = os.path.join(cls.save_dir, "val")

        for path in [cls.save_model_dir, cls.save_log_dir, cls.save_cfg_dir, cls.save_val_dir]:
            # if os.path.exists(path):
            #     shutil.rmtree(path)
            os.makedirs(path, exist_ok=True)

        from corekit.utils.path import WORK_DIR  #

        cls.config_path = os.path.join(cls.save_cfg_dir, 'modelmgt.json')
        cls.agents_path = os.path.join(WORK_DIR, 'hiklayout', 'group_info.json')  # os.getenv('work_dir', '')
        cls.rule_engine_path = os.path.join(WORK_DIR, 'hiklayout', 'rule_engine_info.json')  # os.getenv('work_dir', '')

    @classmethod
    def visual_summary(cls, data, step, name="scalars"):
        """
        将训练过程中的数据写入log文件，方便可视化
        Args:
            data: 字典格式，每个step的保存数据，例如{'loss': 0.2, 'avg loss': 0.1}
            step: 整型，训练步数
            name: str
        """
        event_path = os.path.join(cls.save_log_dir, name)
        for k, v in data.items():
            path = os.path.join(event_path, k, "events.bin")
            value = [step, v, float(time.time())]
            if os.path.exists(path):
                save_data = np.fromfile(path, dtype=float)
            else:
                if not os.path.exists(os.path.dirname(path)):
                    os.makedirs(os.path.dirname(path))
                save_data = np.array([])
            save_data = np.insert(save_data, len(save_data), value)
            save_data.tofile(path)

        if cls.writer:
            cls.writer.add_scalars(name, data, step)

    @classmethod
    def write_group_info(cls, info: dict):
        """新需求, 将联盟训练模式下的训练池信息写到特定位置, 以实现在网页显示信息.
        :param info: dict, 联赛训练的训练池信息
        :return:
        """
        with open(cls.agents_path, 'w', encoding='utf-8') as fp:
            json.dump(info, fp, indent=4, ensure_ascii=False)

    @classmethod
    def write_rule_info(cls, info: dict):
        """新需求, 将规则引擎信息写到特定位置，以实现在网页显示信息.
        :param info: dict, 规则引擎返回信息
        :return:
        """
        with open(cls.rule_engine_path, 'w', encoding='utf-8') as fp:
            json.dump(info, fp, indent=4, ensure_ascii=False)

    @classmethod
    def write_train_metrics(cls, info_dict):
        """
        写训练集评估信息
        :param info_dict: 训练集评估信息字典，例如 {"acc": 0.90, "loss": 0.02}
        :return:
        """
        cls.config['version_info']['train_metrics'] = info_dict
        with open(cls.config_path, 'w+') as config_file:
            json.dump(cls.config, config_file, indent=4, ensure_ascii=False)

    @classmethod
    def write_val_metrics(cls, info_dict):
        """
        写验证集评估信息
        :param info_dict: 验证集评估信息字典，例如 {"acc": 0.90, "loss": 0.02}
        :return:
        """
        cls.config['version_info']['val_metrics'] = info_dict
        with open(cls.config_path, 'w+') as config_file:
            json.dump(cls.config, config_file)

    @classmethod
    def _encrypt_model(cls, path, encrypt_key):
        cls.config['version_info']['encrypt_key'] = encrypt_key
        with open(cls.config_path, 'w+') as config_file:
            json.dump(cls.config, config_file)
        files = os.listdir(path)
        for file in files:
            file = os.path.abspath(os.path.join(path, file))
            if os.path.isfile(file) and ".enc" not in file:
                os.system("openssl enc -aes-128-ecb -in %s -out %s -nosalt -k %s -md md5" % (
                    file, file + ".enc", encrypt_key))
                os.remove(file)
            elif os.path.isdir(file):
                cls._encrypt_model(file, encrypt_key)

    @classmethod
    def save_model(cls, path):
        """
        将模型和配置文件保存到对应目录下
        :param path: 模型文件路径
        :return:
        """
        if not isinstance(path, str):
            raise ValueError("模型路径不是一个字符串(string).")

        if not os.path.exists(path):
            raise FileNotFoundError("模型路径: %s 不存在." % path)

        if os.path.isfile(path):
            shutil.copy(path, cls.save_model_dir)
        elif os.path.isdir(path):
            copyfiles(path, cls.save_model_dir)
        else:
            pass
        if cls.encrypt_key:
            cls._encrypt_model(os.path.join(cls.save_model_dir), cls.encrypt_key)

    @classmethod
    def save_config_file(cls, path):
        """
        将配置参数文件保存到对应目录下
        :param path: 配置文件路径, 字符串或者列表
        :return:
        """
        if not isinstance(path, str) and not isinstance(path, list):
            raise Exception("配置文件路径不是字符串或者字符串的列表.")

        if isinstance(path, str):
            path = [path]

        for path_ in path:
            if not os.path.exists(path_) or not os.path.isfile(path_):
                raise Exception("路径 {} 不存在或者不是文件.".format(path_))
            temp_path = os.path.join(cls.save_cfg_dir, os.path.basename(path_))
            shutil.copyfile(path_, temp_path)

    @classmethod
    def save_classes_info(cls, classes, formats="csv"):
        """
        保存类别信息
        :param classes: 映射信息字典，例如 {"dog": 0, "cat": 1} 或者路径
        :param formats: str, 默认格式
        :return:
        """
        if not isinstance(classes, str) and not isinstance(classes, dict):
            raise Exception("类别信息不是一个路径字符串或者字典.")

        if formats not in ["csv", "txt"]:
            raise Exception("路径信息保存模式只能为 'csv' 或者 'txt'.")

        if isinstance(classes, dict):
            classes_file = os.path.join(
                cls.save_cfg_dir, 'classes.%s' % formats.lower())
            with open(classes_file, 'w+') as f:
                for k, v in classes.items():
                    f.write("{},{}\n".format(k, v))

        else:
            dst_path = os.path.join(
                cls.save_cfg_dir, os.path.basename(classes))
            shutil.copyfile(classes, dst_path)

    @classmethod
    def search_file(cls, path, label):
        """
        搜索指定路径下带标签的文件
        """
        results = search_files(path, label)
        return results

    @classmethod
    def download_dataset(cls, path):
        """
        下载数据到指定路径
        """
        print("dataset starts to download...")
        if os.path.exists(path):
            # raise Exception("%s 已经存在！"%path)
            shutil.rmtree(path)
        shutil.copytree(cls.dataset_dir, path)
        print("dataset downloads done.")

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

    @classmethod
    def get_pretrained_model(cls):
        return cls.pretrain_path

    @classmethod
    def save_model_graph(cls, img_path):
        """
        将配置参数文件保存到对应目录下
        :param img_path: 配置文件
        :return:
        """
        try:
            shutil.copy(img_path, cls.save_model_dir)
        except Exception as e:
            print(e)
            pass

        save_img_path = os.path.join(
            cls.save_model_dir, os.path.basename(img_path))
        cls.config['version_info']['model_graph'] = save_img_path
        with open(cls.config_path, 'w+') as config_file:
            json.dump(cls.config, config_file)

    @classmethod
    def save_infer_json(cls, info_dict):
        """
        将字典信息转换成推理所需的配置文件按
        :param info_dict: 映射信息字典，例如 {"anchor": 100, "label": 2}
        :return:
        """
        if "anchor" in info_dict and "label" in info_dict and "type" in info_dict and "inputSize" in info_dict:
            with open(os.path.join(cls.save_model_dir, 'obj_det_cfg.json'), 'w+') as json_file:
                data = json.dumps(info_dict)
                json_file.write(data)
        else:
            print("格式不正确！请检查是否有anchor, label, type, inputSize等字段")

    @classmethod
    def init_metric(cls, save_dir=None):
        """
        将字典信息转换成推理所需的配置文件
        """
        if not save_dir:
            cls.eval_result_dir = os.getenv("save_dir")
        else:
            cls.eval_result_dir = save_dir
        cls.eval_metrics_dir = os.path.join(cls.eval_result_dir, "metrics")
        cls.eval_detail_dir = os.path.join(cls.eval_result_dir, "results")
        for path in [cls.eval_metrics_dir, cls.eval_detail_dir]:
            # if os.path.exists(path):
            #     shutil.rmtree(path)
            os.makedirs(path, exist_ok=True)
