# -*- coding: utf-8 -*-
import numpy as np

from collections import namedtuple
NoduleBox = namedtuple('NoduleBox',
                       ['z', 'y', 'x', 'diameter', 'is_pos', 'probability'])


def nodule_raw2standard(nodule_boxes, raw_spacing, standard_spacing, start=(0, 0, 0)):
    new_nodule_boxes = []

    for nodule_box in nodule_boxes:
        z, y, x, diameter, is_pos, probability = nodule_box

        z = z * raw_spacing[0] / standard_spacing[0] - start[0]
        y = y * raw_spacing[1] / standard_spacing[1] - start[1]
        x = x * raw_spacing[2] / standard_spacing[2] - start[2]
        new_nodule_boxes.append(
            NoduleBox(z, y, x, diameter, is_pos, probability))

    return new_nodule_boxes


def nodule_standard2raw(nodule_boxes, standard_spacing, raw_spacing, start=(0, 0, 0)):
    new_nodule_boxes = []

    for nodule_box in nodule_boxes:
        z, y, x, diameter, is_pos, probability = nodule_box

        z = (z + start[0]) * standard_spacing[0] / raw_spacing[0]
        y = (y + start[1]) * standard_spacing[1] / raw_spacing[1]
        x = (x + start[2]) * standard_spacing[2] / raw_spacing[2]
        new_nodule_boxes.append(
            NoduleBox(z, y, x, diameter, is_pos, probability))

    return new_nodule_boxes


def iou2d(box_a, box_b):
    box_a = np.asarray(box_a)
    box_b = np.asarray(box_b)

    a_start, a_end = box_a[0:2] - box_a[2:4] / 2, box_a[0:2] + box_a[2:4] / 2
    b_start, b_end = box_b[0:2] - box_b[2:4] / 2, box_b[0:2] + box_b[2:4] / 2

    y_overlap = max(0, min(a_end[0], b_end[0]) - max(a_start[0], b_start[0]))
    x_overlap = max(0, min(a_end[1], b_end[1]) - max(a_start[1], b_start[1]))

    intersection = y_overlap * x_overlap
    union = box_a[2] * box_a[3] + box_b[2] * box_b[3] - intersection

    return 1.0 * intersection / union


def iou3d(box_a, box_b, spacing):
    # 半径需要在z, y, x各轴上换算成像素值
    r_a = 0.5 * box_a.diameter / spacing
    r_b = 0.5 * box_b.diameter / spacing

    # starting index in each dimension
    z_a_s, y_a_s, x_a_s = box_a.z - r_a[0], box_a.y - r_a[1], box_a.x - r_a[2]
    z_b_s, y_b_s, x_b_s = box_b.z - r_b[0], box_b.y - r_b[1], box_b.x - r_b[2]
    # ending index in each dimension
    z_a_e, y_a_e, x_a_e = box_a.z + r_a[0], box_a.y + r_a[1], box_a.x + r_a[2]
    z_b_e, y_b_e, x_b_e = box_b.z + r_b[0], box_b.y + r_b[1], box_b.x + r_b[2]

    z_overlap = max(0, min(z_a_e, z_b_e) - max(z_a_s, z_b_s))
    y_overlap = max(0, min(y_a_e, y_b_e) - max(y_a_s, y_b_s))
    x_overlap = max(0, min(x_a_e, x_b_e) - max(x_a_s, x_b_s))

    intersection = z_overlap * y_overlap * x_overlap
    union = 8 * r_a[0] * r_a[1] * r_a[2] + 8 * r_b[0] * r_b[1] * r_b[2] - intersection

    return 1.0 * intersection / union


def nms(nodule_boxes, spacing, iou_thred):
    if nodule_boxes is None or len(nodule_boxes) == 0:
        return nodule_boxes

    nodule_boxes_nms = []

    for nodule_box in nodule_boxes:
        overlap = False

        for box in nodule_boxes_nms:
            if iou3d(box, nodule_box, spacing) >= iou_thred:
                overlap = True
                break

        if not overlap:
            nodule_boxes_nms.append(nodule_box)

    return nodule_boxes_nms