# coding = utf-8
import os
# import tensorboardX
import argparse
import xml.etree.ElementTree as ET
import time
import json
import shutil
import numpy as np


class RLAlgBase(object):
    """
        算法集成的功能类，只是基础类
    Args:

    """
    @classmethod
    def init(cls):
        cls.save_dir = os.getenv("save_dir")
        cls._init_model_package()

        cls.dataset_dir = os.getenv("dataset_path")
        # tensorbordX 写入
        # cls.writer = tensorboardX.SummaryWriter(cls.save_log_dir)
        cls._init_config_json()

    @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 _init_model_package(cls):
        """
        初始化模型包的格式
        """
        '''
        模型包包含`log`, `model`，`val`和`config`四部分，其中,
        `log`中包含用于可视化模型结构及结果的`tensorboard`文件，
        `model`包含训练输出的模型文件，
        `val`包含模型评估对每个输出数据的预测结果，
        `config`包含该模型版本的基础信息和在其创建的模型评估、模型压缩、模型转换任务。
        '''
        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")
        # 由于创建训练时会初始化旧版本的模型包，但是在该代码编写时模型包已经改了，
        # 所以需要将旧版本的删除，后续可以删除下面的代码
        if os.path.exists(cls.save_dir):
            shutil.rmtree(cls.save_dir, ignore_errors=True)
        for path in [cls.save_dir, cls.save_model_dir, cls.save_log_dir, cls.save_cfg_dir]:
            os.makedirs(path, exist_ok=True)

    @classmethod
    def _init_config_json(cls):
        """
        初始化配置文件夹里的文件，由于是json文件，
        因此先用一个变量记住，后续一次性保存
        """
        cls.config = {
            "user": "",
            "preset": "",
            "name": "",
            "usage": "",
            "engine": "",
            "scene": "",
            "model_type": "",
            "description": ""
        }
