
import os
import time
import warnings

from mmcv import Config

from hikframe.dl.transfer.hikdet.mmdet.apis import train_detector
from hikframe.dl.transfer.hikdet.mmdet.datasets import build_dataset, build_dataloader
from hikframe.dl.transfer.hikdet.mmdet.models import build_detector

from hikframe.dl.transfer.hikdet.mmdet.apis import (inference_detector, init_detector, imwrite_result_pyplot, single_gpu_test)
from mmcv.runner import load_checkpoint
from mmcv.parallel import MMDataParallel
import numpy as np
import shutil
import json

class HikTransfer():
    def __init__(self):
        pass

    def build_dataset(self, datapath):
        self.datasets = [build_dataset(datapath)]
        return self.datasets

    def load_config(self,path_config):
        self.path_cfg = path_config
        self.cfg = Config.fromfile(path_config)

        if 'runner' not in self.cfg:
            self.cfg.runner = {
                'type': 'EpochBasedRunner',
                'max_epochs': self.cfg.total_epochs
            }
            warnings.warn(
                'config is now expected to have a `runner` section, '
                'please set `runner` in your config.', UserWarning)

        self.cfg.gpu_ids = range(1)
        self.cfg.seed = 0

        return self.cfg

    def build_model(self,cfg):
        self.model = build_detector(
            self.cfg.model,
            train_cfg=self.cfg.get('train_cfg'),
            test_cfg=self.cfg.get('test_cfg'))
        return self.model

    def train_model(self,model, datasets, cfg, validate=True):
        timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
        model.CLASSES = datasets[0].CLASSES
        train_detector(
            model,
            datasets,
            cfg,
            distributed=False,
            validate=validate,
            timestamp=timestamp,
            meta=None)

    def test_model(self,model, dataset, cfg):

        data_loader = build_dataloader(
            dataset[0],
            samples_per_gpu=1,
            workers_per_gpu=1,
            dist=False,
            shuffle=False)

        model = MMDataParallel(model, device_ids=[0])
        outputs = single_gpu_test(model, data_loader)
        eval_kwargs = cfg.get('evaluation', {}).copy()
        # hard-code way to remove EvalHook args
        for key in [
            'interval', 'tmpdir', 'start', 'gpu_collect', 'save_best',
            'rule'
        ]:
            eval_kwargs.pop(key, None)
        # eval_kwargs.update(dict(metric=args.eval, **kwargs))
        print(dataset[0].evaluate(outputs, **eval_kwargs))


    def init_model(self,model, checkpoint):
        load_checkpoint(model, checkpoint, map_location='cpu')
        return model

    def inference_model(self, cfg_path, path_checkpoint, img_path, score_thr=0.5, isShow=False,  out_img_path='result_show.jpg', device="cuda:0"):
        # build the model from a config file and a checkpoint file
        model = init_detector(cfg_path, path_checkpoint, device=device)
        # test a single image
        result = inference_detector(model, img_path)
        # result[0] = result[0][ind,:]
        # ind = np.where(result[0][:, 4] >= score_thr)[0]
        for i in range(len(result)):
            ind = np.where(result[i][:,4]>=score_thr)[0]
            result[i] = result[i][ind,:]

        if isShow==True:
            # show the results
            imwrite_result_pyplot(model, img_path, result, score_thr=score_thr, out_file=out_img_path)

        return result

    def download_datastet(self, data_path):
        if not os.path.exists(data_path):
            os.makedirs(data_path)

        shutil.copytree(self.cfg.data_root, data_path)

    def save_config_file(self, path_cfg_save):
        name_file = os.path.splitext(os.path.basename(self.path_cfg))[0]
        self.cfg.dump(path_cfg_save)
        # shutil.copy(os.path.join(self.cfg.work_dir, os.path.basename(self.path_cfg)), path_cfg_save)
        # shutil.copy(os.path.join(self.cfg.work_dir+'/'+name_file+'/', os.path.basename(self.path_cfg)), path_cfg_save)

    def save_from_file(self, path_model_save):
        name_file = os.path.splitext(os.path.basename(self.path_cfg))[0]
        shutil.copy(self.cfg.work_dir+'/latest.pth', path_model_save)
        # shutil.copy(self.cfg.work_dir + '/' + name_file + '/latest.pth', path_model_save)

    def visual_summary(self):
        files = os.listdir(self.cfg.work_dir)

        file_latest_json = ''
        for file in files:
            ind = file.find('log.json')
            if ind> -1:
                if file_latest_json == '':
                    file_latest_json = file
                else:
                    int_file_latest = int(file_latest_json.split('.log.json')[0])
                    int_file_tmp = int(file.split('.log.json')[0])
                    if int_file_latest<int_file_tmp:
                        file_latest_json =  file

        loss_dict = {'epoch': [], 'iter': [], 'lr': [], 'loss_cls': [], 'loss_conf': [], 'loss_xy': [], 'loss_wh': [], 'loss': []}
        with open(self.cfg.work_dir + '/'+file_latest_json, 'r', encoding='utf-8') as fd:
            lines = fd.readlines()
            for i in range(1, len(lines)):
                data = json.loads(lines[i])
                if 'loss' in data.keys() and data['mode']=='train':
                    for key in loss_dict.keys():
                        loss_dict[key].append(data[key])
        return  loss_dict

        # with open(path_log, 'r') as fd:
        #     lines = fd.readlines()
        #
        # loss = {}
        # for line in lines:
        #     ind_val = line.find('mmdet - INFO - Epoch')
        #     if ind_val>-1:
        #         ind_1 = line.find('[')
        #         ind_2 = line.find(']')
        #         key = int(line[ind_1:ind_2+1])
        #
        #         ind_loss = line.find('loss: ')
        #         ind_end = line.find(', grad_n')
        #
        #         if key in loss:
        #             loss[key].append(float(line[ind_loss, ind_end]))
        # keys = loss.keys()
        # keys.sort()
        # loss_ave = []
        # for i in keys:
        #     loss_ave.append(np.mean(loss[i]))
        #
        # return {'loss':loss, 'avg_loss':loss_ave}

    # def show_result(self, cfg_path, path_checkpoint, img_path,  out_img_path, score_thr=0.5, device="cuda:0"):
    #     # build the model from a config file and a checkpoint file
    #     model = init_detector(cfg_path, path_checkpoint, device=device)
    #     # test a single image
    #     result = inference_detector(model, img_path)
    #
    #     # show the results
    #     imwrite_result_pyplot(model, img_path, result, score_thr=score_thr, out_file=out_img_path)

def demo_hikTransfer_test_mAP():
    # 初始化对象
    hiktransfer = HikTransfer()

    path_cfg = 'configs/yolo/rr_yolov4_cspd53_416_coco.py'
    # 导入config配置
    cfg = hiktransfer.load_config(path_cfg)

    # 可修改config配置
    # 如：迁移模型类别、目标类别、保存模型名称、结果路径等
    cfg.model.backbone.type = 'RRCSPDarknet53'
    # cfg.model.class_list = ['cat', 'dog', 'sheep']
    cfg.model.bbox_head.num_classes = 2
    cfg.modelprefix = 'transfered_model_yolo'
    cfg.work_dir = 'training_result'
    # 迁移学习优化策略及超参，如学习率、训练轮次
    cfg.optimizer.lr = 0.0001
    max_epochs = 12
    cfg.total_epochs = max_epochs
    cfg.runner.max_epochs = max_epochs
    cfg.seed = np.random.randint(100)
    cfg.gpu_ids = [0]

    # 构建自定义模型
    cfg.data.test.ann_file = './data/tmp_list_test.txt'
    dataset = hiktransfer.build_dataset(cfg.data.test)

    # 构建模型
    model = hiktransfer.build_model(cfg)

    # 导入离线模型包
    checkpoint = './work_dirs/exps/rr_yolov4_cspd53_416_coco2/epoch_49.pth'
    model = hiktransfer.init_model(model, checkpoint)

    # 执行批量图片性能测试
    hiktransfer.test_model(model, dataset, cfg)


def demo_hikTransfer_train():
    # 初始化对象
    hiktransfer = HikTransfer()


    path_cfg = 'configs/yolo/rr_yolov4_cspd53_416_coco.py'
    # 导入config配置
    cfg = hiktransfer.load_config(path_cfg)

    # 可修改config配置
    # 如：迁移模型类别、目标类别、保存模型名称、结果路径等
    cfg.model.backbone.type = 'RRCSPDarknet53'
    # cfg.model.class_list = ['cat', 'dog', 'sheep']
    cfg.model.bbox_head.num_classes = 2
    cfg.modelprefix = 'transfered_model_yolo'
    cfg.work_dir = 'training_result'
    # 迁移学习优化策略及超参，如学习率、训练轮次
    cfg.optimizer.lr = 0.0001
    max_epochs = 12
    cfg.total_epochs = max_epochs
    cfg.runner.max_epochs = max_epochs
    cfg.seed = np.random.randint(100)
    cfg.gpu_ids = [0]

    # 构建自定义模型
    cfg.data.train.ann_file = './data/tmp_list_train.txt'
    dataset = hiktransfer.build_dataset(cfg.data.train)


    # 构建模型
    model = hiktransfer.build_model(cfg)

    # 迁移学习训练
    cfg.load_from = './pretrained_models/rr_yolov4_pretrained.pth'
    # hiktransfer.train_model(model, dataset, cfg, validate=True)


    hiktransfer.save_config_file('./tmp_cfg.py')
    hiktransfer.save_from_file('./tmp.pth')
    loss_dict = hiktransfer.visual_summary()

def demo_hikTransfer_inference():
    # 初始化对象
    hiktransfer = HikTransfer()

    # 模型推理
    img_path = './demo/vehicle_person1.jpg'
    model_path = './pretrained_models/rr_yolov4_pretrained.pth'
    path_cfg = './configs/yolo/rr_yolov4_cspd53_416_coco.py'
    result = hiktransfer.inference_model(path_cfg, model_path, img_path, score_thr=0.5)
    print(result)
    #
    # # # 模型推理结果可视化
    out_file = 'visual_result.jpg'
    hiktransfer.inference_model(path_cfg, model_path, img_path, score_thr=0.5, isShow=True,  out_img_path=out_file)


if __name__ == '__main__':
    demo_hikTransfer_train()
    # demo_hikTransfer_test_mAP()
    # demo_hikTransfer_inference()
