import logging
import json

from mmcv.ops import get_compiling_cuda_version, get_compiler_version
import mmdet
# Check Pytorch installation
import torch
from mmdet.utils import collect_env, get_root_logger
from mmcv.ops import get_compiling_cuda_version, get_compiler_version
import mmdet
from mmdet.apis import inference_detector, init_detector
from mmcv import Config
from mmdet.apis import set_random_seed
import mmcv
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector
import os.path as osp
import shutil
# from hikframe.dl.base.base import DLAlgBase
from hikframe.dl.base.train import HikModel
import os
import pdb
import numpy as np
from hikframe.dl.cv.hikdet.core.hook import MyTensorboardLoggerHook
import hikframe
from hikframe.dl.utils import visual_model
from mmcv.runner import get_dist_info, init_dist

# config放进来然后可以load和下载
# 批量推理图片，结果+画图
# 训练
# 测试


class HikDet(HikModel):
    @classmethod
    def init(cls, save_dir=None, data_dir=None):
        print('HIKFRAME VERSION:', hikframe.__version__)
        super().__init__(save_dir, data_dir)
        # Check Pytorch installation
        print(torch.__version__, torch.cuda.is_available())
        # Check MMDetection installation
        print(mmdet.__version__)
        # Check mmcv installation
        print(get_compiling_cuda_version())
        print(get_compiler_version())

    @classmethod
    def download_dataset(cls, path):
        return super().download_dataset(path)

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

    @classmethod
    def get_dev_dir(cls):
        cls.tmp_save_dir = osp.join(cls.save_dir, "train_results")
        # try:
        #     if osp.exists(cls.tmp_save_dir):
        #         shutil.rmtree(cls.tmp_save_dir)
        # except Exception as e:
        #     print(e)
        os.makedirs(cls.tmp_save_dir, exist_ok=True)
        return cls.tmp_save_dir

    @classmethod
    def get_pretrained_model(cls):
        pretrained = os.getenv('pretrain')
        pth_path = None
        if pretrained == "true":
            model_dir = os.getenv('restore_model_path')
            pth_path = cls._get_ckpt(model_dir)
            if pth_path:
                return pth_path
            print('[!!!] no pretrained model got.')
        return pth_path

    @classmethod
    def _get_ckpt(cls, model_dir):
        pth_path = None
        if os.path.isfile(model_dir) and model_dir.endswith('.pth'):
            pth_path = model_dir
            return pth_path
        for root, dirs, files in os.walk(model_dir):
            # 优先找model文件夹
            if 'model' in dirs:
                model_dir = osp.join(root, 'model')
                pth_path = cls._find_pth(model_dir)
                if pth_path:
                    return pth_path
            # 继续寻找
            for dir_name in dirs:
                model_dir = osp.join(root, dir_name)
                pth_path = cls._find_pth(model_dir)
                if pth_path:
                    return pth_path
        return pth_path

    @classmethod
    def _find_pth(cls, model_dir):
        pth_path = None
        ckpt_list = [x for x in os.listdir(
            model_dir) if x.endswith('.pth')]
        if not ckpt_list:
            return pth_path
        if 'latest.pth' in ckpt_list:
            ckpt_name = 'latest.pth'
            pth_path = osp.join(model_dir, ckpt_name)
        else:
            try:
                ckpt_list.sort(key=lambda x: int(
                    osp.splitext(x)[0].split('_')[-1]))
            except:
                ckpt_list.sort(key=lambda x: int(
                    osp.splitext(x)[0]))
            if ckpt_list:
                ckpt_name = ckpt_list[-1]
                pth_path = osp.join(model_dir, ckpt_name)
        return pth_path

    @classmethod
    def init_model(cls, cfg, datasets):
        # Build the detector
        model = build_detector(
            cfg.model, train_cfg=cfg.get('train_cfg'), test_cfg=cfg.get('test_cfg'))
        # Add an attribute for visualization convenience
        model.CLASSES = datasets[0].CLASSES
        # Create work_dir
        mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
        return model

    @classmethod
    def restore_model(cls, cfg, pth_path=None, resume_flag=False):
        if not pth_path:
            pth_path = HikDet.get_pretrained_model()
        if not pth_path or not os.path.isfile(pth_path):
            print('[ERROR] model restore failed, pth_path must direct to a pth file!, but input is {}'.format(
                pth_path))
            return cfg
        # 若环境变量没传resume，则以接口传入为准
        resume = os.getenv('resume')
        if resume:
            resume = True if resume == 'true' else False
        else:
            resume = resume_flag
        if pth_path:
            if resume:
                cfg.resume_from = pth_path
            else:
                cfg.load_from = pth_path
        return cfg

    @classmethod
    def config_from_file(cls, config_file):
        cfg = Config.fromfile(config_file)
        # Set seed thus the results are more re producible
        tbhook = None
        for hook in cfg.log_config.hooks:
            if hook['type'] == 'TensorboardLoggerHook':
                tbhook = hook
                break
        if tbhook:
            cfg.log_config.hooks.remove(tbhook)
        log_dir = cls.save_log_dir if hasattr(
            cls, 'save_log_dir') else cfg.work_dir
        # cfg.log_config.hooks.append(
        #     dict(type='TensorboardLoggerHook', log_dir=log_dir))
        cfg.log_config.hooks.append(
            dict(type='MyTensorboardLoggerHook', log_dir=log_dir))
        cfg.seed = 0
        set_random_seed(0, deterministic=False)
        cfg.gpu_ids = range(1)
        # set distribute mode, if use multi-gpu this must be set as true
        cls.distribute = True if os.getenv("distribute") == 'true' else False
        if os.getenv("distribute") == 'true':
            init_dist('pytorch', **cfg.dist_params)
        return cfg

    @classmethod
    def save_config(cls, cfg):
        cfg_file = osp.join(cls.save_cfg_dir, 'config.py')
        cfg.dump(cfg_file)
        print('config.py is saved under {}'.format(cfg_file))

    @classmethod
    def init_logger(cls, save_dir=''):
        if not save_dir:
            save_dir = cls.tmp_save_dir
        log_file = osp.join(cls.save_cfg_dir, 'logger.log')
        logger = get_root_logger(log_file=log_file, log_level=logging.INFO)
        # log env info
        env_info_dict = collect_env()
        env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()])
        dash_line = '-' * 60 + '\n'
        logger.info('Environment info:\n' + dash_line + env_info + '\n' +
                    dash_line)
        return logger

    @classmethod
    def save_model_graph(cls, cfg, file_dir):
        datasets = [build_dataset(cfg.data.train)]
        model = HikDet.init_model(cfg, datasets)
        inputs = torch.rand(1, 3, 256, 256)
        pic_path = visual_model(
            model, inputs, framework='hikdet', file_dir=file_dir)
        super().save_model_graph(pic_path)
        return pic_path

    @classmethod
    def train(cls, cfg):
        datasets = [build_dataset(cfg.data.train)]
        model = HikDet.init_model(cfg, datasets)
        distributed = cls.distribute if hasattr(cls, 'distribute') else False
        train_detector(model, datasets, cfg,
                       distributed=distributed, validate=True)

        cls._cosntruct_model_pkg(cls.tmp_save_dir)

        # get metric_result
        metric_result = {}
        for file in os.listdir(cls.save_cfg_dir):
            try:
                if file.endswith('.log.json'):
                    file_path = os.path.join(cls.save_cfg_dir, file)
                    metric_result = cls._parser_log(file_path, cfg)
                    if metric_result:
                        break
            except:
                metric_result = {}
                pass

        super().write_train_metrics(metric_result)
        super().write_val_metrics(metric_result)
        return metric_result

    @classmethod
    def _parser_log(cls, json_file, cfg):
        # ['proposal', 'bbox']
        metric_list = cfg.evaluation.metric
        metric_result = {}
        with open(json_file, 'r') as f:
            for line in f.readlines()[::-1]:
                info = json.loads(line)
                if info["mode"] != "val":
                    continue
                for k, v in info.items():
                    if 'AR' in k:
                        metric_result[k] = float(v)
                    elif '_mAP_copypaste' in k:
                        ret = v.split(' ')
                        if len(ret) == 6:
                            name = ['mAP', 'mAP_50', 'mAP_75',
                                    'mAP_s', 'mAP_m', 'mAP_l']
                            for n, value in zip(name, ret):
                                # print(n, value)
                                metric_result[n] = float(value)
                    # user defined
                    for metric in metric_list:
                        if k == metric:
                            metric_result[k] = float(v)
                if metric_result:
                    return metric_result
        return metric_result

    @ classmethod
    def infer(cls, model, img_dir, out_dir='', score_threshold=0.8, show_result=False):
        """
        input:
        img_dir: img_file_path or img_dir(contain images)
        return:
        dict, like{'xx.jpg':[(label_text, score, xmin, ymin, xmax, ymax), ...], ...}
        will save predict.txt under out_dir and visual img under out_dir/visual_result
        """
        rets = {}
        # class_names = cfg.data.train.classes
        class_names = model.CLASSES
        print('class_names:', class_names)

        if show_result:
            if not out_dir:
                out_dir = '/out_infer'
            os.makedirs(out_dir, exist_ok=True)
            img_out_dir = osp.join(out_dir, 'visual_result')
            os.makedirs(img_out_dir, exist_ok=True)
        # fit single img file path
        if osp.isfile(img_dir):
            imgs = [img_dir]
        else:
            imgs = os.listdir(img_dir)
        for img in imgs:
            # if not img.endswith(('jpg', 'png', 'jpeg', 'png', 'bmp')):
            #     continue
            if osp.exists(img):
                imgpath = img
            else:
                imgpath = osp.join(img_dir, img)
            det_rets = inference_detector(model, imgpath)
            img_name = imgpath
            rets[img_name] = []
            ret = HikDet._det_result_parser(
                class_names, det_rets, score_threshold)
            rets[img_name] = ret
            if show_result:
                out_file = osp.join(
                    img_out_dir, '{}.jpg'.format(os.path.splitext(os.path.basename(img))[0]))
                print(out_file)
                model.show_result(imgpath, det_rets, score_thr=score_threshold,
                                  show=False, out_file=out_file)
        # print(f'infer result saved under {prediction_path}.')
        if show_result:
            print(f'visual result saved under {img_out_dir}.')

        return rets

    @ classmethod
    def _det_result_parser(cls, class_names, det, score_threshold):
        """
        return: list, like [(label_text, score, xmin, ymin, xmax, ymax), ...]
        """
        ret = []
        bbox_result, segm_result = det, None
        bboxes = np.vstack(bbox_result)
        labels = [
            np.full(bbox.shape[0], i, dtype=np.int32)
            for i, bbox in enumerate(bbox_result)
        ]
        labels = np.concatenate(labels)
        if score_threshold >= 0:
            assert bboxes.shape[1] == 5
            scores = bboxes[:, -1]
            inds = scores > score_threshold
            scores = scores[inds]
            bboxes = bboxes[inds, :]
            labels = labels[inds]
        for i, (score, bbox, label) in enumerate(zip(scores, bboxes, labels)):
            bbox_int = bbox.astype(np.int32)
            xmin, ymin, xmax, ymax = bbox_int[:4]
            label_text = class_names[
                label] if class_names is not None else f'class {label}'
            ret.append((label_text, score, xmin, ymin, xmax, ymax))
        return ret

    @ classmethod
    def build_model(cls, cfg, ckpt_path=''):
        if not ckpt_path:
            return
        ckpt_file = ckpt_path
        if osp.isdir(ckpt_path):
            ckpt_file = cls._get_ckpt(ckpt_path)
        if not (ckpt_file and ckpt_file.endswith('.pth')):
            return
        try:
            if torch.cuda.is_available():
                model = init_detector(cfg, ckpt_file, device='cuda:0')
            else:
                model = init_detector(cfg, ckpt_file, device='cpu')
        except Exception as e:
            print(f'ERROR: model load error, {e}')
            model = None
        return model

    @ classmethod
    def video_infer(cls):
        # todo
        pass

    @ classmethod
    def evaluate(cls):
        # todo
        pass

    @ classmethod
    def _cosntruct_model_pkg(cls, save_dir):
        assert save_dir and osp.exists(save_dir)
        for root, dirs, files in os.walk(save_dir):
            for f in files:
                # copy model
                if f.endswith('.pth'):
                    # shutil.copy(osp.join(root, f), os.path.join(
                    #     cls.save_model_dir, f))
                    shutil.move(osp.join(root, f), os.path.join(
                        cls.save_model_dir, f))
                # copy config and log
                elif f.endswith(('.log', '.json', '.py')):
                    shutil.move(osp.join(root, f),
                                osp.join(cls.save_cfg_dir, f))
                elif f.startswith('events.out'):
                    shutil.move(osp.join(root, f),
                                osp.join(cls.save_log_dir, f))
                else:
                    pass
        # try:
        #     shutil.rmtree(cls.tmp_save_dir)
        # except Exception as e:
        #     print(e)
        print('model package constructed success.')
