from collections import OrderedDict

from mmcv.utils import print_log

from hikframe.dl.transfer.hikdet.mmdet.core import eval_map, eval_recalls
from .builder import DATASETS
from .custom import CustomDataset

import os.path as osp
import xml.etree.ElementTree as ET

import mmcv
import numpy as np
from PIL import Image


@DATASETS.register_module()
class CJLDataset(CustomDataset):

    CLASSES = ('vehicle', 'person')

    def __init__(self, min_size=None, **kwargs):
        assert self.CLASSES or kwargs.get(
            'classes', None), 'CLASSES in `XMLDataset` can not be None.'
        super(CJLDataset, self).__init__(**kwargs)
        self.cat2label = {cat: i for i, cat in enumerate(self.CLASSES)}
        self.min_size = min_size

    def load_annotations(self, ann_file):
        """Load annotation from XML style ann_file.

        Args:
            ann_file (str): Path of XML file.

        Returns:
            list[dict]: Annotation info from jpg file.
        """

        data_infos = []
        img_ids = mmcv.list_from_file(ann_file)
        for img_id in img_ids:
            filename = img_id
            # xml_subnanme = osp.splitext(filename)[0] +'.xml'
            # xml_path = osp.join(self.img_prefix, xml_subnanme.replace('JPEGImages', 'Annotations'))
            # tree = ET.parse(xml_path)
            # root = tree.getroot()
            # size = root.find('size')
            # if size is not None:
            #     width = int(size.find('width').text)
            #     height = int(size.find('height').text)
            # else:

            img_path = osp.join(self.img_prefix, filename)
            img = Image.open(img_path)
            width, height = img.size
            data_infos.append(
                dict(id=img_id, filename=filename, width=width, height=height))

        return data_infos

    def _filter_imgs(self, min_size=32):
        """Filter images too small or without annotation."""
        valid_inds = []
        for i, img_info in enumerate(self.data_infos):
            if min(img_info['width'], img_info['height']) < min_size:
                continue
            # if self.filter_empty_gt:
            #     filename = img_info['filename']
            #     xml_subnanme = osp.splitext(filename)[0] + '.xml'
            #     xml_path = osp.join(self.img_prefix,xml_subnanme.replace('JPEGImages', 'Annotations'))
            #
            #     tree = ET.parse(xml_path)
            #     root = tree.getroot()
            #     for obj in root.findall('object'):
            #         name = obj.find('name').text
            #         if name in self.CLASSES:
            #             valid_inds.append(i)
            #             break
            # else:
            #     valid_inds.append(i)

            valid_inds.append(i)
        return valid_inds

    def get_ann_info(self, idx):
        """Get annotation from labels txt file by index.

        Args:
            idx (int): Index of data.

        Returns:
            dict: Annotation info of specified index.
        """

        filename = self.data_infos[idx]['filename']
        width = self.data_infos[idx]['width']
        height = self.data_infos[idx]['height']

        txt_subnanme = osp.splitext(filename)[0] + '.txt'
        txt_path = osp.join(self.img_prefix,txt_subnanme.replace('JPEGImages', 'labels'))
        data = np.loadtxt(txt_path)

        if len(data.shape) == 1:
            data = np.expand_dims(data, axis =0)
        # print('data shape: ', data.shape)
        bboxes = []
        labels = []
        bboxes_ignore = []
        labels_ignore = []
        # for obj in root.findall('object'):
        for i in range(data.shape[0]):
            # print(txt_path)
            # print(data.shape)
            label = data[i, 0]
            # TODO: check whether it is necessary to use int
            # Coordinates may be float type
            bbox = [
                int(width *(data[i, 1] -data[i, 3]/2.0)),
                int(height *(data[i, 2] -data[i, 4]/2.0)),
                int(width *(data[i, 1] +data[i, 3]/2.0)),
                int(height *(data[i, 2] +data[i, 4]/2.0))
            ]
            difficult = False
            ignore = False
            if self.min_size:
                assert not self.test_mode
                w = bbox[2] - bbox[0]
                h = bbox[3] - bbox[1]
                if w < self.min_size or h < self.min_size:
                    ignore = True
            if difficult or ignore:
                bboxes_ignore.append(bbox)
                labels_ignore.append(label)
            else:
                bboxes.append(bbox)
                labels.append(label)
        if not bboxes:
            bboxes = np.zeros((0, 4))
            labels = np.zeros((0, ))
        else:
            bboxes = np.array(bboxes, ndmin=2) - 1
            labels = np.array(labels)
        if not bboxes_ignore:
            bboxes_ignore = np.zeros((0, 4))
            labels_ignore = np.zeros((0, ))
        else:
            bboxes_ignore = np.array(bboxes_ignore, ndmin=2) - 1
            labels_ignore = np.array(labels_ignore)
        ann = dict(
            bboxes=bboxes.astype(np.float32),
            labels=labels.astype(np.int64),
            bboxes_ignore=bboxes_ignore.astype(np.float32),
            labels_ignore=labels_ignore.astype(np.int64))
        # print('ann---------------, ', ann)
        return ann

    def get_cat_ids(self, idx):
        """Get category ids in XML file by index.

        Args:
            idx (int): Index of data.

        Returns:
            list[int]: All categories in the image of specified index.
        """

        cat_ids = []
        filename = self.data_infos[idx]['filename']
        txt_subnanme = osp.splitext(filename)[0] + '.txt'
        txt_path = osp.join(self.img_prefix, txt_subnanme.replace('JPEGImages', 'labels'))
        data = np.loadtxt(txt_path)
        if len(data.shape) == 1:
            data = np.newaxis(data)

        for i in range(data.shape[0]):
            label = data[i, 0]
            cat_ids.append(label)
        # print('cat_ids---------------, ', cat_ids)

        return cat_ids

    def evaluate(self,
                 results,
                 metric='mAP',
                 logger=None,
                 proposal_nums=(100, 300, 1000),
                 iou_thr=0.5,
                 scale_ranges=None):
        """Evaluate in VOC protocol.

        Args:
            results (list[list | tuple]): Testing results of the dataset.
            metric (str | list[str]): Metrics to be evaluated. Options are
                'mAP', 'recall'.
            logger (logging.Logger | str, optional): Logger used for printing
                related information during evaluation. Default: None.
            proposal_nums (Sequence[int]): Proposal number used for evaluating
                recalls, such as recall@100, recall@1000.
                Default: (100, 300, 1000).
            iou_thr (float | list[float]): IoU threshold. Default: 0.5.
            scale_ranges (list[tuple], optional): Scale ranges for evaluating
                mAP. If not specified, all bounding boxes would be included in
                evaluation. Default: None.

        Returns:
            dict[str, float]: AP/recall metrics.
        """

        if not isinstance(metric, str):
            assert len(metric) == 1
            metric = metric[0]
        allowed_metrics = ['mAP', 'recall']
        if metric not in allowed_metrics:
            raise KeyError(f'metric {metric} is not supported')
        annotations = [self.get_ann_info(i) for i in range(len(self))]
        eval_results = OrderedDict()
        iou_thrs = [iou_thr] if isinstance(iou_thr, float) else iou_thr
        if metric == 'mAP':
            assert isinstance(iou_thrs, list)

            ds_name = self.CLASSES
            mean_aps = []
            for iou_thr in iou_thrs:
                print_log(f'\n{"-" * 15}iou_thr: {iou_thr}{"-" * 15}')
                mean_ap, _ = eval_map(
                    results,
                    annotations,
                    scale_ranges=None,
                    iou_thr=iou_thr,
                    dataset=ds_name,
                    logger=logger)
                mean_aps.append(mean_ap)
                eval_results[f'AP{int(iou_thr * 100):02d}'] = round(mean_ap, 3)
            eval_results['mAP'] = sum(mean_aps) / len(mean_aps)
        elif metric == 'recall':
            gt_bboxes = [ann['bboxes'] for ann in annotations]
            recalls = eval_recalls(
                gt_bboxes, results, proposal_nums, iou_thr, logger=logger)
            for i, num in enumerate(proposal_nums):
                for j, iou in enumerate(iou_thr):
                    eval_results[f'recall@{num}@{iou}'] = recalls[i, j]
            if recalls.shape[1] > 1:
                ar = recalls.mean(axis=1)
                for i, num in enumerate(proposal_nums):
                    eval_results[f'AR@{num}'] = ar[i]
        return eval_results
