import torch
import numpy as np
import math

def bbox_bep(box1, box2, xywh=True, eps=1e-7, bep1 = True):
    """
    Calculates bottom edge proximity between two boxes
    
    Input shapes are box1(1,4) to box2(n,4)
    
    Implementation of bep2 from 
        Are object detection assessment criteria ready for maritime computer vision?
    """

    # Get the coordinates of bounding boxes
    if xywh:  # transform from xywh to xyxy
        (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
        w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
        b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
        b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
    else:  # x1, y1, x2, y2 = box1
        b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)
        b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1)
        w1, h1 = b1_x2 - b1_x1, (b1_y2 - b1_y1).clamp(eps)
        w2, h2 = b2_x2 - b2_x1, (b2_y2 - b2_y1).clamp(eps)

    # Bottom edge distance (absolute value)
    # xb = torch.abs(b2_x2 - b1_x1)
    xb = torch.min(b2_x2-b1_x1, b1_x2-b2_x1)
    xa = w2 - xb
    xc = w1 - xb
    ybe = torch.abs(b2_y2 - b1_y2)

    X2 = xb/(xb+xa)
    Y2 = 1-ybe/h2

    X1 = xb/(xb+xa+xc+eps)
    Y1 = 1-ybe/(torch.max(h2,h1)+eps)

    bep = X1*Y1 if bep1 else X2*Y2

    return bep

def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
    """
    Calculates IoU, GIoU, DIoU, or CIoU between two boxes, supporting xywh/xyxy formats.

    Input shapes are box1(1,4) to box2(n,4).
    """

    # Get the coordinates of bounding boxes
    if xywh:  # transform from xywh to xyxy
        (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
        w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
        b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
        b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
    else:  # x1, y1, x2, y2 = box1
        b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)
        b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1)
        w1, h1 = b1_x2 - b1_x1, (b1_y2 - b1_y1).clamp(eps)
        w2, h2 = b2_x2 - b2_x1, (b2_y2 - b2_y1).clamp(eps)

    # Intersection area
    inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp(0) * (
        b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)
    ).clamp(0)

    # Union Area
    union = w1 * h1 + w2 * h2 - inter + eps

    # IoU
    iou = inter / union
    if CIoU or DIoU or GIoU:
        cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1)  # convex (smallest enclosing box) width
        ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1)  # convex height
        if CIoU or DIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
            c2 = cw**2 + ch**2 + eps  # convex diagonal squared
            rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center dist ** 2
            if CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
                v = (4 / math.pi**2) * (torch.atan(w2 / h2) - torch.atan(w1 / h1)).pow(2)
                with torch.no_grad():
                    alpha = v / (v - iou + (1 + eps))
                return iou - (rho2 / c2 + v * alpha)  # CIoU
            return iou - rho2 / c2  # DIoU
        c_area = cw * ch + eps  # convex area
        return iou - (c_area - union) / c_area  # GIoU https://arxiv.org/pdf/1902.09630.pdf
    return iou  # IoU

class BoxMetrics:
    # Updated version of https://github.com/kaanakan/object_detection_confusion_matrix
    def __init__(self):
        self.preds_tm = []
        self.target_tm = []
        self.bottom_x = []
        self.bottom_y = []
        self.widths = []
        self.heights = []
        self.ious = []
        self.beps = []

    def add_batch(self, preds, target):
        """
        Return intersection-over-union (Jaccard index) of boxes.
        Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
        Arguments:
            detections torch(Array[N, 6]), x1, y1, x2, y2, conf, class 
            labels torch(Array[M, 5]), class, x1, y1, x2, y2
        Returns:
            None, updates confusion matrix accordingly
        """
        self.preds_tm.extend(preds)
        self.target_tm.extend(target)
    
    def compute(self):
        """
        Computes bbox iou, bep and location/size statistics
        """

        for i in range(len(self.target_tm)):
            target_batch_boxes = self.target_tm[i][:, 1:]
            pred_batch_boxes = self.preds_tm[i][:, :4]

            if pred_batch_boxes.shape[0] == 0:
                continue

            if target_batch_boxes.shape[0] == 0:
                continue

            for t_box in target_batch_boxes:
                    iou = bbox_iou(t_box.unsqueeze(0), pred_batch_boxes, xywh=False)
                    bep = bbox_bep(t_box.unsqueeze(0), pred_batch_boxes, xywh=False)

                    matches = pred_batch_boxes[iou.squeeze(1) > 0.1]

                    bep = bep[iou > 0]
                    iou = iou[iou > 0]
                    # if any iou value is 0 or less, raise error
                    if torch.any(iou <= 0):
                        raise ValueError("IoU values must be greater than 0.")
                    #same for bep
                    if torch.any(bep <= 0):
                        print(t_box.unsqueeze(0))
                        print(pred_batch_boxes)
                        print(bep)
                        print(iou)
                        raise ValueError("BEP values must be greater than 0.")
                    
                    self.ious.extend(iou.tolist())
                    self.beps.extend(bep.tolist())
                    
                    for match in matches:
                        t_xc = (match[0].item()+match[2].item())/2
                        p_xc = (t_box[0].item()+t_box[2].item())/2
                        t_w = t_box[2].item()-t_box[0].item()
                        p_w = match[2].item()-match[0].item()
                        t_h = t_box[3].item()-t_box[1].item()
                        p_h = match[3].item()-match[1].item()

                        self.bottom_x.append(p_xc - t_xc)
                        self.bottom_y.append(match[3].item()-t_box[3].item())
                        self.widths.append(p_w-t_w)
                        self.heights.append(p_h-t_h)

        return {"iou_mean": np.mean(self.ious),
                "bep_mean": np.mean(self.beps),
                "bottom_x_std": np.std(self.bottom_x),
                "bottom_y_std": np.std(self.bottom_y),
                "widths_std": np.std(self.widths),
                "heights_std": np.std(self.heights),
                "bottom_x_mean": np.mean(self.bottom_x),
                "bottom_y_mean": np.mean(self.bottom_y),
                "widths_mean": np.mean(self.widths),
                "heights_mean": np.mean(self.heights)}