import os
import shutil

try:
    import cv2
except:
    print('WARNING: cv2 is not installed.')
    pass
import json
from corekit.hikframe.dl.utils.filesys import search_files
from corekit.hikframe.dl.utils.env import print_envs


class Evaluation(object):
    """
    评估类
    """

    @classmethod
    def __init__(cls, save_dir=None, dataset_dir=None, model_dir=None):
        """
        初始化函数
        """
        # 显示传入的env参数
        cls.envs = ['gpu_num', 'code_dir',
                    'dataset_dir', 'save_dir', 'model_dir']
        print_envs(cls.envs)

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

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

        cls.model_dir = model_dir if model_dir else os.getenv("model_dir")
        if cls.model_dir is None:
            raise Exception("未指定模型路径. 设置环境变量'model_dir'或者输入初始化参数 'model_dir'")

        # 从配置文件中读取密钥
        config_path = os.path.join(cls.model_dir, 'config/modelmgt.json')
        cls.encrypt_key = None
        with open(config_path) as f:
            model_config = json.load(f)
            try:
                cls.encrypt_key = model_config["version_info"]["encrypt_key"]
            except:
                pass

        # 初始化评估报告路径
        cls._init_report_package()

        # 初始化指标信息
        cls._init_metrics()

    @classmethod
    def _init_report_package(cls):
        """
        初始化评估报告包的格式

        报告包含`metrics`, `predict.txt` 和`results`三部分，其中,
        `metrics`中包含评估指标信息、展示图片以及表格信息，
        `predict.txt`中是每个样本的预测信息，
        `results`包含针对每个样本的模型评估结果。
        """
        # 一级目录
        cls.metrics_path = os.path.join(cls.save_dir, "metrics")
        cls.predict_path = os.path.join(cls.save_dir, "predict.txt")
        cls.results_path = os.path.join(cls.save_dir, "results")

        # 二级目录
        cls.images_path = os.path.join(cls.metrics_path, "images")
        cls.metrics_path = os.path.join(cls.metrics_path, "metrics.json")
        cls.wrong_path = os.path.join(cls.results_path, "wrong")
        cls.right_path = os.path.join(cls.results_path, "right")

        for path in [cls.images_path, cls.wrong_path, cls.right_path]:
            os.makedirs(path, exist_ok=True)

    @classmethod
    def _init_metrics(cls):
        """
        初始化指标信息
        """
        cls.metrics = {
            "index": {
                "brief": {},  # 总体指标
                "detail": []  # 每个类别指标
            },
            "chart": []  # 图表展示
        }

    @classmethod
    def _update_metrics(cls):
        if not os.path.isfile(cls.metrics_path):
            return
        with open(cls.metrics_path, 'r') as f:
            cls.metrics = json.load(f)
            return

    @classmethod
    def brief_index(cls, names, values):
        """
        更新总体指标
        parmas: names 列表，指标的名字
               values 列表，指标名字对应的数值
        """
        if not isinstance(names, list) or not isinstance(values, list):
            raise Exception("输入的参数 'names' 和 'values' 并不是列表.")

        if len(names) != len(values):
            raise Exception("输入的参数 'names' 和 'values' 列表长度不一样.")

        cls._update_metrics()

        for name, value in zip(names, values):
            cls.metrics["index"]["brief"][name] = value

        with open(cls.metrics_path, "w") as f:
            json.dump(cls.metrics, f)

    @classmethod
    def detail_index(cls, mark, names, values, image_path):
        """
        详细指标
        parmas: mark 标记，详细指标的标记/类别，唯一标识
                names 列表，指标的名字
                values 列表，指标名字对应的数值
                image_path 可视化指标图路径
        """
        if not isinstance(names, list) or not isinstance(values, list):
            raise Exception("输入的参数 'names' 和 'values' 并不是列表.")

        if len(names) != len(values):
            raise Exception("输入的参数 'names' 和 'values' 列表长度不一样.")

        if not isinstance(mark, str):
            raise Exception("详细指标的标识 'mark' 不是字符串.")

        data = dict()
        for name, value in zip(names, values):
            data[name] = value

        dst_path = None
        if os.path.exists(image_path):
            _, fmt = os.path.splitext(image_path)
            dst_path = os.path.join(cls.images_path, "%s%s" % (mark, fmt))
            shutil.copyfile(image_path, dst_path)

        data["image_path"] = dst_path

        cls._update_metrics()
        cls.metrics["index"]["detail"].append({mark: data})
        with open(cls.metrics_path, "w") as f:
            json.dump(cls.metrics, f)

    @classmethod
    def chart(cls, name, data, mode):
        """
        图表信息
        params: name 图表名字
                data 图表数据
                mode 图表格式 bar/pie
        """
        if not isinstance(name, str):
            raise Exception("表名 'name' 不是字符串.")

        if not isinstance(data, list):
            raise Exception(
                "表数据不是列表. 柱状图和饼图数据格式应该为[{name: test, value: 4}, ...]")

        if mode in ["bar", "pie"]:
            # 柱状图和饼图
            for item in data:
                if not isinstance(item, dict) or not isinstance(item["name"], str) or not isinstance(item["value"],
                                                                                                     (int, float)):
                    raise Exception(
                        "单个数据不是字典或者数据不符合格式. 柱状图和饼图数据格式应该为[{'name': string, 'value': float/int}, ...] ")

                item["name"] = str(item["name"])

        else:
            raise Exception("只支持 bar（柱状图）和 pie（饼图）两种模式.")

        cls._update_metrics()
        cls.metrics["chart"].append({
            "name": name,
            "mode": mode,
            "data": data
        })
        with open(cls.metrics_path, "w") as f:
            json.dump(cls.metrics, f)

    @classmethod
    def judge(cls, path, ground_truth, predictions, mode, is_right=True, save_name=None):
        """
        存储单个样本预测对错
        params: path 字符串
                ground_truth 列表或者字符串，样本标签真值
                predictions 列表，预测结果
                is_right 布尔型，是否预测正确
                save_name 字符串，保存json文件名
        """
        data = dict()
        if not isinstance(path, str):
            raise Exception("样本路径 path 需要为字符串.")
        if not os.path.exists(path):
            raise Exception("样本路径 path 不存在.")
        data["path"] = path

        # 判断是否需要图片路径
        _, fmt = os.path.splitext(path)
        if fmt in [".jpg", ".JPG", ".png", ".PNG", ".bmp", ".BMP", ".JPEG", ".jpeg"]:
            img = cv2.imread(path)
            size = img.shape
            data["size"] = {
                "width": str(size[1]), "height": str(
                    size[0]), "depth": str(size[2])
            }

        # 真值
        if not isinstance(ground_truth, list):
            raise Exception(
                "样本标签 ground_truth 需要为列表. 分类时为['car'],...]；检测时为[['car', 0.1, 0.1, 0.9. 0.9],...]")

        trans_gts = []
        for item in ground_truth:
            if mode == "detect":
                trans_gts.append({
                    "label": item[0],
                    "bbox": item[1:]
                })
            else:
                trans_gts.append({
                    "label": item[0]
                })

        data["gt"] = trans_gts

        # 预测值
        if not isinstance(predictions, list):
            raise Exception(
                "样本预测值 predictions 需要列表. 分类时为[['car', 0.9],...]；检测时为[['car', 0.1, 0.1, 0.9. 0.9],...]")

        trans_pred = []
        for item in predictions:
            if mode == "detect":
                trans_pred.append({
                    "label": item[0],
                    "bbox": item[2:],
                    "score": item[1]
                })
            else:
                trans_pred.append({
                    "label": item[0],
                    "score": item[1]
                })

        data["predictions"] = trans_pred

        # 存成json
        save_path = cls.right_path if is_right else cls.wrong_path
        save_name = save_name if save_name else os.path.splitext(
            os.path.basename(path))[0]
        os.makedirs(save_path, exist_ok=True)
        with open("%s.json" % os.path.join(save_path, save_name), "w") as f:
            json.dump(data, f)

        # 写入predict
        if mode == "classify":
            line = "{} {}\n".format(path, predictions)
        elif mode == "detect":
            pred = []
            for item in predictions:
                pred += item
            line = "{} {}\n".format(
                path, " ".join([str(item) for item in pred]))
        else:
            raise Exception("只支持classify(分类)和detect(检测)两种模式.")

        with open(cls.predict_path, "a+") as f:
            f.writelines([line])

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