import numpy as np

import math

# Error Messages

METAFILE_NOTFOUND_ERR = "Metadata file {} could not be parsed! Exception: {}!"

LOAD_FAIL_MSG = "Failed to load image: {}\nException: {}"

# Constants

IMG_PAD_TOKEN = "<PAD>"

VOXEL_SPACING = (0.703125, 0.703125, 2.5)

CENSORING_DIST = {

    "0": 0.9851933387778276,

    "1": 0.9748326267838882,

    "2": 0.9659936099043455,

    "3": 0.9587266943913851,

    "4": 0.952360797355704,

    "5": 0.9461857341570268,

}

def order_slices(img_paths, slice_locations):

    sorted_ids = np.argsort(slice_locations)

    sorted_img_paths = np.array(img_paths)[sorted_ids].tolist()

    sorted_slice_locs = np.sort(slice_locations).tolist()

    return sorted_img_paths, sorted_slice_locs

def assign_splits(meta, args):

    for idx in range(len(meta)):

        meta[idx]["split"] = np.random.choice(

            ["train", "dev", "test"], p=args.split_probs

        )

def get_scaled_annotation_mask(additional, args, scale_annotation=True):

    """

    Construct bounding box masks for annotations

    Args:

        - additional['image_annotations']: list of dicts { 'x', 'y', 'width', 'height' }, where bounding box coordinates are scaled [0,1].

        - args

    Returns:

        - mask of same size as input image, filled in where bounding box was drawn. If additional['image_annotations'] = None, return empty mask. Values correspond to how much of a pixel lies inside the bounding box, as a fraction of the bounding box's area

    """

    H, W = args.img_size

    mask = np.zeros((H, W))

    if additional["image_annotations"] is None:

        return mask

    for annotation in additional["image_annotations"]:

        single_mask = np.zeros((H, W))

        x_left, y_top = annotation["x"] * W, annotation["y"] * H

        x_right, y_bottom = (

            x_left + annotation["width"] * W,

            y_top + annotation["height"] * H,

        )

        # pixels completely inside bounding box

        x_quant_left, y_quant_top = math.ceil(x_left), math.ceil(y_top)

        x_quant_right, y_quant_bottom = math.floor(x_right), math.floor(y_bottom)

        # excess area along edges

        dx_left = x_quant_left - x_left

        dx_right = x_right - x_quant_right

        dy_top = y_quant_top - y_top

        dy_bottom = y_bottom - y_quant_bottom

        # fill in corners first in case they are over-written later by greater true intersection

        # corners

        single_mask[math.floor(y_top), math.floor(x_left)] = dx_left * dy_top

        single_mask[math.floor(y_top), x_quant_right] = dx_right * dy_top

        single_mask[y_quant_bottom, math.floor(x_left)] = dx_left * dy_bottom

        single_mask[y_quant_bottom, x_quant_right] = dx_right * dy_bottom

        # edges

        single_mask[y_quant_top:y_quant_bottom, math.floor(x_left)] = dx_left

        single_mask[y_quant_top:y_quant_bottom, x_quant_right] = dx_right

        single_mask[math.floor(y_top), x_quant_left:x_quant_right] = dy_top

        single_mask[y_quant_bottom, x_quant_left:x_quant_right] = dy_bottom

        # completely inside

        single_mask[y_quant_top:y_quant_bottom, x_quant_left:x_quant_right] = 1

        # in case there are multiple boxes, add masks and divide by total later

        mask += single_mask

    if scale_annotation:

        mask /= mask.sum()

    return mask

def get_scaled_annotation_area(sample, args):

    """

    no_box = [{'width': 0, 'height': 0}]

    if sample['series'] in self.annotations_metadata:

        # total area of bounding boxes in

        areas_per_slice = [ [ box['width']*box['height'] for box in self.annotations_metadata[ sample['series'] ].get( os.path.splitext(os.path.basename(path))[0], no_box ) ] for path in sample['paths'] ]

        return np.array( [ np.sum(areas) for areas in areas_per_slice] )

    else:

        return np.array([ 0  for _ in sample['paths'] ])

    """

    areas = []

    for additional in sample["annotations"]:

        mask = get_scaled_annotation_mask(additional, args, scale_annotation=False)

        areas.append(mask.sum() / (mask.shape[0] * mask.shape[1]))

    return np.array(areas)