import math
import pickle
import numpy as np
import os 
from scipy.io import loadmat


def get_width(boxes, offset=0):
    return boxes[:, 2] - boxes[:, 0] + offset


def get_height(boxes, offset=0):
    return boxes[:, 3] - boxes[:, 1] + offset


def get_size(boxes, offset=0):
    return get_width(boxes, offset) * get_height(boxes, offset)

def overlap(boxes_a, boxes_b, mode='iou', is_aligned=False, offset=0, eps=1e-6):
    # boxes_a: (n, 4)
    # boxes_b: (m, 4)
    # return: (n, m)
    if is_aligned:
        assert boxes_a.shape[0] == boxes_b.shape[0]
    if boxes_a.shape[0] * boxes_b.shape[0] == 0:
        if is_aligned:
            return np.zeros((boxes_a.shape[0], ))
        else:
            return np.zeros((boxes_a.shape[0], boxes_b.shape[0]))
    if is_aligned:
        lt = np.maximum(boxes_a[:, :2], boxes_b[:, :2])  # (n, 2)
        rb = np.minimum(boxes_a[:, 2:4], boxes_b[:, 2:4])  # (n, 2)
        wh = np.maximum(rb - lt + offset, 0)  # (n, 2)
        overlap = wh[:, 0] * wh[:, 1]  # (n,)
        if mode is None:
            return overlap
        if mode in ('iou', 'ioa', 'giou', 'diou', 'ciou'):
            area_a = get_size(boxes_a, offset=offset)  # (n,)
        if mode in ('iou', 'iob', 'giou', 'diou', 'ciou'):
            area_b = get_size(boxes_b, offset=offset)  # (n,)
        if mode in ('iou', 'giou', 'diou', 'ciou'):  # intersection over union
            union = area_a + area_b - overlap  # (n,)
        elif mode == 'ioa':  # intersection over a
            union = area_a  # (n,)
        elif mode == 'iob':  # intersection over b
            union = area_b  # (n,)
        else:
            assert False
        if mode == 'giou' or mode == 'diou' or mode == 'ciou':
            enclosed_lt = np.minimum(boxes_a[:, :2], boxes_b[:, :2])  # (n, 2)
            enclosed_rb = np.maximum(boxes_a[:, 2:4], boxes_b[:, 2:4])  # (n, 2)
            enclosed_wh = np.maximum(enclosed_rb - enclosed_lt + offset, 0)  # (n, 2)
            enclosed_area = enclosed_wh[:, 0] * enclosed_wh[:, 1]  # (n,)

            if  mode == 'diou' or mode == 'ciou':
                c2 = enclosed_wh[:,  0] ** 2 +  enclosed_wh[:,  1] ** 2 + eps # convex diagonal squared 
                rho2 = ((boxes_b[:, 0] + boxes_b[:, 2] - boxes_a[:, 0] - boxes_a[:,  2]) ** 2 + (boxes_b[:, 1] + boxes_b[ :, 3] - boxes_a[:, 1] - boxes_a[:, 3]) ** 2) / 4 # center distance squared
                if mode == 'ciou':
                    w1, h1 = boxes_a[:, 2] - boxes_a[:, 0], boxes_a[:, 3] - boxes_a[:, 1]
                    w2, h2 = boxes_b[:, 2] - boxes_b[:, 0], boxes_b[:, 3] - boxes_b[:, 1]    
    else:
        lt = np.maximum(boxes_a[:, None, :2], boxes_b[None, :, :2])  # (n, m, 2)
        rb = np.minimum(boxes_a[:, None, 2:4], boxes_b[None, :, 2:4])  # (n, m, 2)
        wh = np.maximum(rb - lt + offset, 0)  # (n, m, 2)
        overlap = wh[:, :, 0] * wh[:, :, 1]  # (n, m)
        if mode is None:
            return overlap
        if mode in ('iou', 'ioa', 'giou', 'diou', 'ciou', 'siou'):
            area_a = get_size(boxes_a, offset=offset)  # (n, )
        if mode in ('iou', 'iob', 'giou', 'diou', 'ciou', 'siou'):
            area_b = get_size(boxes_b, offset=offset)  # (m, )
        if mode in ('iou', 'giou', 'diou', 'ciou', 'siou'):  # intersection over union
            union = area_a[:, None] + area_b[None, :] - overlap  # (n, m)
        elif mode == 'ioa':  # intersection over a
            union = area_a[:, None]  # (n, m)
        elif mode == 'iob':  # intersection over b
            union = area_b[None, :]  # (n, m)
        else:
            assert False
        if mode == 'giou' or mode == 'diou' or mode == 'ciou' or mode == 'siou':
            enclosed_lt = np.minimum(boxes_a[:, None, :2], boxes_b[None, :, :2])  # (n, m, 2)
            enclosed_rb = np.maximum(boxes_a[:, None, 2:4], boxes_b[None, :, 2:4])  # (n, m, 2)
            enclosed_wh = np.maximum(enclosed_rb - enclosed_lt + offset, 0)  # (n, m, 2)
            enclosed_area = enclosed_wh[:, :, 0] * enclosed_wh[:, :, 1]  # (n, m)
            if  mode == 'diou' or mode == 'ciou':
                c2 = enclosed_wh[:, :, 0] ** 2 +  enclosed_wh[:, :, 1] ** 2 + eps # convex diagonal squared 
                rho2 = ((boxes_b[None, :, 0] + boxes_b[None, :, 2] - boxes_a[:, None, 0] - boxes_a[:, None, 2]) ** 2 + (boxes_b[None, :, 1] + boxes_b[None, :, 3] - boxes_a[:, None, 1] - boxes_a[:, None, 3]) ** 2) / 4 # center distance squared
                if mode == 'ciou':
                    w1, h1 = boxes_a[:, None, 2] - boxes_a[:, None, 0], boxes_a[:, None, 3] - boxes_a[:, None, 1]
                    w2, h2 = boxes_b[None, :, 2] - boxes_b[None, :, 0], boxes_b[None, :, 3] - boxes_b[None, :, 1]

    union = np.maximum(union, eps)
    ios = overlap / union
    if mode == 'giou':
        enclosed_area = np.maximum(enclosed_area, eps)
        giou = ios - (enclosed_area - union) / enclosed_area
        return giou
    elif mode == 'diou':
        diou = ios - rho2 / c2
        return diou
    elif mode == 'ciou':
        v = (4 / math.pi ** 2) * math.pow(math.atan(w2 / (h2 + eps)) - math.atan(w1 / h1), 2)
        alpha = v / (v - ios + (1+ eps))
        ciou = ios - (rho2 / c2 + v * alpha)
        return ciou
    else:
        return ios


def average_precision(recalls, precisions, mode='area'):
    """Calculate average precision (for single or multiple scales).

    Args:
        recalls (ndarray): shape (num_scales, num_dets) or (num_dets, )
        precisions (ndarray): shape (num_scales, num_dets) or (num_dets, )
        mode (str): 'area' or '11points', 'area' means calculating the area
            under precision-recall curve, '11points' means calculating
            the average precision of recalls at [0, 0.1, ..., 1]

    Returns:
        float or ndarray: calculated average precision
    """
    no_scale = False
    if recalls.ndim == 1:
        no_scale = True
        recalls = recalls[np.newaxis, :]
        precisions = precisions[np.newaxis, :]
    assert recalls.shape == precisions.shape and recalls.ndim == 2
    num_scales = recalls.shape[0]
    ap = np.zeros(num_scales, dtype=np.float32)
    if mode == 'area':
        zeros = np.zeros((num_scales, 1), dtype=recalls.dtype)
        ones = np.ones((num_scales, 1), dtype=recalls.dtype)
        mrec = np.hstack((zeros, recalls, ones))
        mpre = np.hstack((zeros, precisions, zeros))
        for i in range(mpre.shape[1] - 1, 0, -1):
            mpre[:, i - 1] = np.maximum(mpre[:, i - 1], mpre[:, i])
        for i in range(num_scales):
            ind = np.where(mrec[i, 1:] != mrec[i, :-1])[0]
            ap[i] = np.sum((mrec[i, ind + 1] - mrec[i, ind]) * mpre[i, ind + 1])
    elif mode == '11points':
        for i in range(num_scales):
            for thr in np.arange(0, 1 + 1e-3, 0.1):
                precs = precisions[i, recalls[i, :] >= thr]
                prec = precs.max() if precs.size > 0 else 0
                ap[i] += prec
            ap /= 11
    else:
        raise ValueError('Unrecognized mode, only "area" and "11points" are supported')
    if no_scale:
        ap = ap[0]
    return ap




def get_gt_boxes(gt_dir):
    """ gt dir: (wider_face_val.mat, wider_easy_val.mat, wider_medium_val.mat, wider_hard_val.mat)"""

    gt_mat = loadmat(os.path.join(gt_dir, 'wider_face_val.mat'))
    hard_mat = loadmat(os.path.join(gt_dir, 'wider_hard_val.mat'))
    medium_mat = loadmat(os.path.join(gt_dir, 'wider_medium_val.mat'))
    easy_mat = loadmat(os.path.join(gt_dir, 'wider_easy_val.mat'))

    facebox_list = gt_mat['face_bbx_list']
    event_list = gt_mat['event_list']
    file_list = gt_mat['file_list']

    hard_gt_list = hard_mat['gt_list']
    medium_gt_list = medium_mat['gt_list']
    easy_gt_list = easy_mat['gt_list']

    return facebox_list, event_list, file_list, hard_gt_list, medium_gt_list, easy_gt_list


def norm_score(pred):
    """ norm score
    pred {key: [[x1,y1,x2,y2,s]]}
    """

    max_score = 0
    min_score = 1

    for _, key in pred.items():
        for _, value in key.items():
            if len(value) == 0:
                continue
            _min = np.min(value[:, -1])
            _max = np.max(value[:, -1])
            max_score = max(_max, max_score)
            min_score = min(_min, min_score)

    diff = max_score - min_score
    for _, key in pred.items():
        for _, value in key.items():
            if len(value) == 0:
                continue
            value[:, -1] = (value[:, -1] - min_score) / diff


def image_eval(pred, gt_boxes, ignore, iou_thresh):
    """ single image evaluation
    pred: Nx5
    gt_boxes: Nx4
    ignore:
    """

    _pred = pred.copy()
    _gt = gt_boxes.copy()
    pred_recall = np.zeros(_pred.shape[0])
    recall_list = np.zeros(_gt.shape[0])
    proposal_list = np.ones(_pred.shape[0])

    _pred[:, 2] = _pred[:, 2] + _pred[:, 0]
    _pred[:, 3] = _pred[:, 3] + _pred[:, 1]
    _gt[:, 2] = _gt[:, 2] + _gt[:, 0]
    _gt[:, 3] = _gt[:, 3] + _gt[:, 1]

    overlaps = overlap(_pred[:, :4], _gt)

    for h in range(_pred.shape[0]):
        gt_overlap = overlaps[h]
        max_overlap, max_idx = gt_overlap.max(), gt_overlap.argmax()
        if max_overlap >= iou_thresh:
            if ignore[max_idx] == 0:
                recall_list[max_idx] = -1
                proposal_list[h] = -1
            elif recall_list[max_idx] == 0:
                recall_list[max_idx] = 1

        r_keep_index = np.where(recall_list == 1)[0]
        pred_recall[h] = len(r_keep_index)
    return pred_recall, proposal_list


def voc_ap(rec, prec):
    """cal mAP in VOC style"""

    # correct AP calculation
    # first append sentinel values at the end
    mrec = np.concatenate(([0.], rec, [1.]))
    mpre = np.concatenate(([0.], prec, [0.]))

    # compute the precision envelope
    for i in range(mpre.size - 1, 0, -1):
        mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

    # to calculate area under PR curve, look for points
    # where X axis (recall) changes value
    i = np.where(mrec[1:] != mrec[:-1])[0]

    # and sum (\Delta recall) * prec
    ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap


def dataset_pr_info(thresh_num, pr_curve, count_face):
    """cal dataset precision recall info"""

    _pr_curve = np.zeros((thresh_num, 2))
    for i in range(thresh_num):
        _pr_curve[i, 0] = pr_curve[i, 1] / pr_curve[i, 0]
        _pr_curve[i, 1] = pr_curve[i, 1] / count_face
    return _pr_curve


def img_pr_info(thresh_num, pred_info, proposal_list, pred_recall):
    """img precision recall info"""

    pr_info = np.zeros((thresh_num, 2)).astype('float')
    for t in range(thresh_num):

        thresh = 1 - (t + 1) / thresh_num
        r_index = np.where(pred_info[:, 4] >= thresh)[0]
        if len(r_index) == 0:
            pr_info[t, 0] = 0
            pr_info[t, 1] = 0
        else:
            r_index = r_index[-1]
            p_index = np.where(proposal_list[:r_index + 1] == 1)[0]
            pr_info[t, 0] = len(p_index)
            pr_info[t, 1] = pred_recall[r_index]
    return pr_info


def parse_pred(pred, pred_names, det_cls=-1):
    """parse model preds
    pred: {'image_name': _ ,
           'pred_boxes': _ ,
           'pred_scores': _}
    """

    pred_box_key = 'pred_boxes'
    pred_score_key = 'pred_scores'
    pred_cls_key = 'pred_clses'
    for pred_name in pred_names:
        if 'boxes' in pred_name:
            pred_box_key = pred_name
        if 'scores' in pred_name:
            pred_score_key = pred_name
        if 'classes' in pred_name:
            pred_cls_key = pred_name
    results = dict()

    suffix_list = ['.jpg', '.pkl', '.raw']
    for _, res in enumerate(pred):
        event, name = res['image_name'].split('/')

        for suffix in suffix_list:
            if name.endswith(suffix):
                name = name.rstrip(suffix)
                break 

        boxes = res[pred_box_key]
        boxes[:, 2] = boxes[:, 2] - boxes[:, 0] + 1
        boxes[:, 3] = boxes[:, 3] - boxes[:, 1] + 1
        scores = res[pred_score_key]
        if det_cls != -1:
            keep = np.where(res[ pred_cls_key ] == det_cls)[0]
            boxes = boxes[keep, ...]
            scores = scores[keep, ...]
        pred_boxes = np.hstack((boxes, scores[:, None])).astype('float64')
        if event in results:
            results[event][name] = pred_boxes
        else:
            results[event] = dict()
            results[event][name] = pred_boxes

    return results