# The draw functions are based from:
# https://github.com/fizyr/keras-retinanet/blob/master/keras_retinanet/utils/visualization.py
# https://github.com/fizyr/keras-maskrcnn/blob/master/keras_maskrcnn/utils/visualization.py

__all__ = [
    "draw_sample",
    "draw_record",
    "draw_pred",
    "draw_bbox",
    "draw_mask",
    "draw_keypoints",
    "draw_label",
    "draw_segmentation_mask",
]

from icevision.imports import *
from icevision.data import *
from icevision.core import *
from icevision.visualize.utils import *
from matplotlib.colors import LinearSegmentedColormap

# This should probably move elsewhere
from PIL import Image, ImageFont, ImageDraw
import PIL

DEFAULT_FONT_PATH = get_default_font()


def draw_sample(
    sample,
    class_map: Optional[ClassMap] = None,
    denormalize_fn: Optional[callable] = None,
    display_label: bool = True,
    display_bbox: bool = True,
    display_score: bool = True,
    display_mask: bool = True,
    display_keypoints: bool = True,
    font_path: Optional[os.PathLike] = DEFAULT_FONT_PATH,
    font_size: Union[int, float] = None,
    label_color: Union[np.array, list, tuple, str] = "#C4C4C4",  # Mild Gray
    label_border_color: Union[np.array, list, tuple, str] = "#020303",  # Black,
    label_thin_border: bool = True,
    label_pad_width_factor: float = 0.02,
    label_pad_height_factor: float = 0.005,
    mask_blend: float = 0.5,
    mask_border_thickness: int = 7,
    color_map: Optional[dict] = None,  # label -> color mapping
    prettify: bool = True,
    prettify_func: Callable = str.capitalize,
    return_as_pil_img=False,
    # Args for plotting specific labels
    exclude_labels: List[str] = [],
    include_only: List[str] = None,
    multiple_classification_spacing_factor: float = 1.05,
    dynamic_font_size_div_factor: float = 20.0,
    include_classification_task_names: bool = True,
    include_instances_task_names: bool = False,
    force_mask_file_reload: bool = False,
) -> Union[np.ndarray, PIL.Image.Image]:
    """
    Selected kwargs:

    * label_color: A <collection> of RGB values or a hex code string that defines
                   the color of all the plotted labels
    * label_border_color: Color of the border around the label
    * label_thin_border: Apply a thin border around the label. If false, applies
                         a thick border. If None, applies no border
    * label_pad_width_factor: Amount of padding to apply relative to the image's width.
                              Applies padding to bbox coords if padding bbox
                              labels else to the top-left of the image for classif labels
    * mask_blend: Degree of transparency of the mask. 1 = opaque, 0 = transparent
    * mask_border_thickness: Degree of thickness of the mask. Must be an odd number
    * color_map: An optional dictionary that maps the label => color-value
    * prettify: Format labels based on `prettify_func`
    * prettify_func: A string -> string processing function
    * return_as_pil_image: If true, returns the sample as a PIL image, else np.array
    * exclude_labels: (Optional) List of labels that you'd like to exclude from being plotted
    * include_only: (Optional) List of labels that must be exclusively plotted. Takes
                    precedence over `exclude_labels` (?)
    """
    img = np.asarray(sample.img).copy()  # HACK
    num_classification_plotted = 0
    # Dynamic font size based on image height
    if font_size is None:
        font_size = sample.img_size.height / dynamic_font_size_div_factor

    if denormalize_fn is not None:
        img = denormalize_fn(img)

    # HACK to visualize segmentation mask
    for task, composite in sample.task_composites.items():
        if task == "segmentation":
            cm = rand_cmap(sample.segmentation.class_map.num_classes, verbose=False)
            if composite.mask_array is None:
                if isinstance(composite.masks[0], RLE):
                    masks = [
                        mask.to_mask(img.shape[1], img.shape[0]).data
                        for mask in composite.masks
                    ]
                elif isinstance(composite.masks[0], EncodedRLEs):
                    masks = composite.masks[0].to_mask(img.shape[0], img.shape[1]).data
                elif isinstance(composite.masks[0], MaskArray):
                    mask = composite.masks[0]
                else:
                    raise ValueError(
                        "Mask has to be of they RLE, EncodedRLEs or MaskArray."
                    )
            else:
                mask = composite.mask_array
            return draw_segmentation_mask(img, mask, cm, display_mask=display_mask)

        # Should break if no ClassMap found in composite.
        #  Should be as the only composite without ClassMap should be
        #  `sample.common`. This is a foundational assumption? #NOTE
        class_map = getattr(composite, "class_map", None)

        if composite.get_component_by_type(ClassificationLabelsRecordComponent):
            x = 0
            y = (
                font_size
                * num_classification_plotted
                * multiple_classification_spacing_factor
            )
            num_classification_plotted += 1
        else:
            x, y = None, None

        # HACK
        if hasattr(composite, "masks"):
            if composite.mask_array is None:
                if isinstance(composite.masks[0], RLE):
                    masks = [
                        mask.to_mask(img.shape[1], img.shape[0]).data
                        for mask in composite.masks
                    ]
                elif isinstance(composite.masks[0], MaskFile):
                    if force_mask_file_reload:
                        logger.warning(
                            "Re-creating masks from files, might results in mismatches if transformations were applied"
                        )
                        masks = [
                            mask.to_mask(img.shape[1], img.shape[0])
                            for mask in composite.masks
                        ]
                    else:
                        logger.warning(
                            "Masks are of type MaskFile but will not be loaded to avoid mismatch with transformed data. Set force_mask_file_reload to True to force mask loading"
                        )
                        masks = []
                elif isinstance(composite.masks[0], EncodedRLEs):
                    masks = composite.masks[0].to_mask(img.shape[0], img.shape[1]).data
                elif isinstance(composite.masks[0], MaskArray):
                    mask = composite.masks[0]
                else:
                    raise ValueError(
                        "Mask has to be of they RLE, EncodedRLEs or MaskArray."
                    )
            else:
                masks = composite.mask_array
        else:
            masks = []

        for label, bbox, mask, keypoints, score in itertools.zip_longest(
            getattr(composite, "labels", []),  # list of strings
            getattr(composite, "bboxes", []),
            masks,
            getattr(composite, "keypoints", []),
            getattr(composite, "scores", []),
        ):
            # random color by default
            color = (np.random.random(3) * 0.6 + 0.4) * 255

            # logic for plotting specific labels only
            # `include_only` > `exclude_labels`
            if not label == []:
                # label_str = (
                #     class_map.get_by_name(label) if class_map is not None else ""
                # )
                if include_only is not None:
                    if not label in include_only:
                        continue
                elif label in exclude_labels:
                    continue

            # if color-map is given and `labels` are predicted
            # then set color accordingly
            if color_map is not None:
                color = as_rgb_tuple(color_map[label])
                color = np.array(color).astype(np.float32)
            if display_mask and mask is not None:
                img = draw_mask(
                    img=img,
                    mask=mask,
                    color=color,
                    blend=mask_blend,
                    border_thickness=mask_border_thickness,
                )
            if display_bbox and bbox is not None:
                img = draw_bbox(img=img, bbox=bbox, color=color)
            if display_keypoints and keypoints is not None:
                img = draw_keypoints(img=img, kps=keypoints, color=color)
            if display_label and label is not None:
                prefix = ""
                if include_classification_task_names:
                    if composite.get_component_by_type(
                        ClassificationLabelsRecordComponent
                    ):
                        prefix = prettify_func(task) + ": "
                if include_instances_task_names:
                    if composite.get_component_by_type(InstancesLabelsRecordComponent):
                        prefix = prettify_func(task) + ": "

                img = draw_label(
                    img=img,
                    label=label,
                    score=score if display_score else None,
                    bbox=bbox,
                    mask=mask,
                    class_map=class_map,
                    color=label_color,
                    border_color=label_border_color,
                    pad_width_factor=label_pad_width_factor,
                    pad_height_factor=label_pad_height_factor,
                    thin_border=label_thin_border,
                    font_size=font_size,
                    font=font_path,
                    prettify=prettify,
                    prettify_func=prettify_func,
                    return_as_pil_img=False,  # should this always be False??
                    prefix=prefix,
                    x=x,
                    y=y,
                )
    if return_as_pil_img:
        # may or may not be a PIL Image based on `display_label`
        return img if isinstance(img, PIL.Image.Image) else PIL.Image.fromarray(img)
    else:
        # will be a `np.ndarray` by default so no need for casting
        return img


def draw_label(
    img: np.ndarray,
    label: Union[int, str],
    score: Optional[float],
    color: Union[np.ndarray, list, tuple],
    border_color: Union[np.ndarray, list, tuple],
    class_map: Optional[ClassMap] = None,
    bbox=None,
    mask=None,
    font: Union[int, os.PathLike, None] = None,
    font_size: Union[int, float] = 12,
    prettify: bool = True,
    prettify_func: Callable = str.capitalize,
    return_as_pil_img=False,
    pad_width_factor=0.02,
    pad_height_factor=0.005,
    thin_border=True,
    x: Optional[int] = None,
    y: Optional[int] = None,
    prefix: str = "",
) -> Union[np.ndarray, PIL.Image.Image]:
    # finds label position based on bbox or mask
    if x is None or y is None:
        # print(f"X: {x}, Y: {y}")
        if bbox is not None:
            x, y, _, _ = bbox.xyxy
        elif mask is not None:
            y, x = np.unravel_index(mask.data.argmax(), mask.data.shape)
        else:
            x, y = 0, 0

    if class_map is not None:
        if isinstance(label, int):
            # TODO: This may never get triggered because we're looping
            # over composite.labels which is a list of strings
            caption = class_map.get_by_id(label)
        else:
            caption = label
    else:
        caption = str(label)
    if prettify:
        # We could introduce a callback here for more complex label renaming
        caption = str(prefix) + str(caption)
        caption = prettify_func(caption)

    # Append label confidence to caption if applicable
    if score is not None:
        if prettify:
            score = f"{score * 100: .2f}%"
        caption = f"{caption}: {score}"

    if not Path(font).exists():
        # PIL throws cryptic errors for wrong filepaths, so let's catch it earlier here
        raise FileNotFoundError(f"{font} file doesn't exist")
    return _draw_label(
        img=img,
        caption=caption,
        x=x,
        y=y,
        color=color,
        border_color=border_color,
        font_path=font,
        font_size=int(font_size),
        return_as_pil_img=return_as_pil_img,
        pad_width_factor=pad_width_factor,
        pad_height_factor=pad_height_factor,
        thin_border=thin_border,
    )


def _draw_label(
    img: np.ndarray,
    caption: str,
    x: int,
    y: int,
    color: Union[np.ndarray, list, tuple],
    border_color: Union[np.ndarray, list, tuple],
    font_path=DEFAULT_FONT_PATH,
    font_size: int = 20,
    return_as_pil_img: bool = False,
    pad_width_factor=0.02,
    pad_height_factor=0.005,
    thin_border=True,
) -> Union[PIL.Image.Image, np.ndarray]:
    """Draw labels on the image"""
    font = PIL.ImageFont.truetype(font_path, size=font_size)
    color = as_rgb_tuple(color)
    border_color = as_rgb_tuple(border_color)

    height, width = img.shape[:2]
    x_pad = height * pad_width_factor
    y_pad = width * pad_height_factor
    x, y = x + x_pad, y + y_pad

    img = PIL.Image.fromarray(img)
    draw = ImageDraw.Draw(img)

    if thin_border is not None:
        # Draw thin / thick border around text
        draw.text(
            (x - 1, y if thin_border else y - 1), caption, font=font, fill=border_color
        )
        draw.text(
            (x + 1, y if thin_border else y - 1), caption, font=font, fill=border_color
        )
        draw.text(
            (x if thin_border else x - 1, y - 1), caption, font=font, fill=border_color
        )
        draw.text(
            (x if thin_border else x + 1, y + 1), caption, font=font, fill=border_color
        )

    # Now draw text over the border
    draw.text((x, y), caption, font=font, fill=color)
    if return_as_pil_img:
        return img
    else:
        return np.array(img)


def draw_record(
    record,
    class_map: Optional[ClassMap] = None,
    display_label: bool = True,
    display_bbox: bool = True,
    display_mask: bool = True,
    display_score: bool = True,
    display_keypoints: bool = True,
    font_path: Optional[os.PathLike] = DEFAULT_FONT_PATH,
    font_size: Union[int, float] = 12,
    label_color: Union[np.array, list, tuple, str] = "#C4C4C4",  # Mild Gray
    mask_blend: float = 0.5,
    mask_border_thickness: int = 7,
    color_map: Optional[dict] = None,  # label -> color mapping
    prettify: bool = True,
    prettify_func: Callable = str.capitalize,
    return_as_pil_img=False,
    # Args for plotting specific labels
    exclude_labels: List[str] = [],
    include_only: List[str] = None,
):
    sample = record.load()
    return draw_sample(
        sample=sample,
        class_map=class_map,
        display_label=display_label,
        display_score=display_score,
        display_bbox=display_bbox,
        display_mask=display_mask,
        display_keypoints=display_keypoints,
        font_path=font_path,
        font_size=font_size,
        label_color=label_color,
        mask_blend=mask_blend,
        mask_border_thickness=mask_border_thickness,
        color_map=color_map,
        prettify=prettify,
        prettify_func=prettify_func,
        return_as_pil_img=return_as_pil_img,
        exclude_labels=exclude_labels,
        include_only=include_only,
    )


def draw_pred(
    pred: Prediction,
    denormalize_fn: Optional[callable] = None,
    display_label: bool = True,
    display_score: bool = True,
    display_bbox: bool = True,
    display_mask: bool = True,
    font_path: Optional[os.PathLike] = DEFAULT_FONT_PATH,
    font_size: Union[int, float] = 12,
    label_color: Union[np.array, list, tuple, str] = "#C4C4C4",  # Mild Gray
    mask_blend: float = 0.5,
    mask_border_thickness: int = 7,
    color_map: Optional[dict] = None,  # label -> color mapping
    prettify: bool = True,
    prettify_func: Callable = str.capitalize,
    return_as_pil_img=False,
    # Args for plotting specific labels
    exclude_labels: List[str] = [],
    include_only: List[str] = None,
):
    return draw_sample(
        sample=pred.pred,
        denormalize_fn=denormalize_fn,
        display_label=display_label,
        display_score=display_score,
        display_bbox=display_bbox,
        display_mask=display_mask,
        font_path=font_path,
        font_size=font_size,
        label_color=label_color,
        mask_blend=mask_blend,
        mask_border_thickness=mask_border_thickness,
        color_map=color_map,
        prettify=prettify,
        prettify_func=prettify_func,
        return_as_pil_img=return_as_pil_img,
        exclude_labels=exclude_labels,
        include_only=include_only,
    )


def draw_bbox(
    img: np.ndarray,
    bbox: BBox,
    color: Tuple[int, int, int],
    gap: bool = True,
):
    """Draws a box on an image with a given color.
    # Arguments
        image     : The image to draw on.
        box       : A list of 4 elements (x1, y1, x2, y2).
        color     : The color of the box.
    """

    # Calculate image dimensions
    dims = sorted(img.shape, reverse=True)
    color = as_rgb_tuple(color)
    img = PIL.Image.fromarray(img)
    draw = PIL.ImageDraw.Draw(img)

    # corner thickness is linearly correlated with the smaller image dimension.
    # We use the smaller image dimension rather than image area so as to avoid
    # overly thick lines for large non-square images prior to transforming
    # images. We set lower and upper bounds for corner thickness.
    min_corner = 1
    max_corner = 15
    corner_thickness = int(0.005 * dims[1] + min_corner)
    if corner_thickness > max_corner:
        corner_thickness = int(max_corner)

    corner_length = int(0.021 * dims[1] + 2.25)

    # inner thickness of bboxes with corners
    inner_thickness = int(1 + 0.0005 * dims[1])

    # bbox thickness of bboxes without corners
    min_bbox = 1
    max_bbox = 8
    bbox_thickness = int(0.0041 * dims[1] - 0.0058)
    if bbox_thickness < min_bbox:
        bbox_thickness = min_bbox
    if bbox_thickness > max_bbox:
        bbox_thickness = int(max_bbox)

    if gap == False:
        xyxy = tuple(np.array(bbox.xyxy, dtype=int))
        draw.rectangle(xyxy, fill=None, outline=color, width=bbox_thickness)
        return np.array(img)

    xmin, ymin, xmax, ymax = tuple(np.array(bbox.xyxy, dtype=int))

    points = [0] * 12
    points[0] = (xmin, ymin + corner_length)
    points[1] = (xmin, ymin)
    points[2] = (xmin + corner_length, ymin)

    points[3] = (xmax - corner_length, ymin)
    points[4] = (xmax, ymin)
    points[5] = (xmax, ymin + corner_length)

    points[6] = (xmax, ymax - corner_length)
    points[7] = (xmax, ymax)
    points[8] = (xmax - corner_length, ymax)

    points[9] = (xmin + corner_length, ymax)
    points[10] = (xmin, ymax)
    points[11] = (xmin, ymax - corner_length)

    if (
        ymax - (ymin + 4 * corner_length) < corner_length
        or xmax - (xmin + 4 * corner_length) < corner_length
    ):
        for i in range(4):
            draw.line(
                xy=(points[i * 3 + 1], points[10 - 3 * i]),
                fill=color,
                width=bbox_thickness,
            )
        for i in range(2):
            draw.line(
                xy=(points[6 * i + 1], points[i * 6 + 4]),
                fill=color,
                width=bbox_thickness,
            )
    else:
        for i in range(2):
            for j in range(2):
                draw.line(
                    xy=(points[i * 6 + j * 4], points[i * 6 + j * 4 + 1]),
                    fill=color,
                    width=corner_thickness,
                    joint=None,
                )
                draw.line(
                    xy=(points[i * 3 + 1], points[10 - 3 * i]),
                    fill=color,
                    width=inner_thickness,
                    joint=None,
                )
        for i in range(2):
            for j in range(2):
                draw.line(
                    xy=(points[i * 6 + j * 2 + 1], points[i * 6 + j * 2 + 2]),
                    fill=color,
                    width=corner_thickness,
                    joint=None,
                )
                draw.line(
                    xy=(points[6 * i + 1], points[i * 6 + 4]),
                    fill=color,
                    width=inner_thickness,
                    joint=None,
                )

    return np.array(img)


def draw_mask(
    img: np.ndarray,
    mask: MaskArray,
    color: Tuple[int, int, int],
    blend: float = 0.5,
    border_thickness: int = 7,
):
    # Border thickness must be an odd integer
    if border_thickness % 2 == 0:
        # TODO: Shall we throw an error, or change the value to the nearest
        # even integer automatically and raise a warning?
        raise ValueError(
            f"`border_thickness` must be an odd number. You entered {border_thickness}"
        )
    img = PIL.Image.fromarray(img)
    w, h = img.size

    mask_idxs = np.where(mask.data)

    # Add alpha with 0 transparency. We draw the mask with border color first
    color = np.append(color, 255)
    mask_arr = np.zeros((h, w, 4), dtype=np.uint8)
    mask_arr[mask_idxs[1:]] = color

    # Now create a second mask and draw the desired color on top of the
    # border mask. If `border_thickness` is 0, this replaces the border mask
    _mask = Image.fromarray(mask_arr)
    _mask = _mask.filter(PIL.ImageFilter.MinFilter(border_thickness))
    _mask_idx = np.where(_mask.convert("L", palette=Image.ADAPTIVE))
    mask_arr[_mask_idx] = np.append(color[:3], blend * 255)

    # Create RGBA PIL mask image
    mask_pil = PIL.Image.fromarray(mask_arr, mode="RGBA")

    # Blend everything keeping the alpha
    # Key concept is that alpha for non-mask pixels are 0 (transparent)
    img.putalpha(255)
    img = PIL.Image.alpha_composite(img, mask_pil)
    return np.array(img)


def draw_segmentation_mask(
    img: np.ndarray,
    mask: MaskArray,
    cmap: LinearSegmentedColormap,
    display_mask: bool = True,
    alpha: float = 0.5,
):
    img = PIL.Image.fromarray(img).convert("RGB")

    if display_mask:
        w, h = img.size
        mask_arr = np.zeros((h, w, 3), dtype=np.uint8)
        mask = mask.data.squeeze()

        assert mask.shape == (h, w), (
            "image and mask size should be the same"
            f"but got image:{(w, h)}; mask: {(mask.shape[::-1])}"
        )

        for class_idx in np.unique(mask):
            mask_idxs = mask == class_idx
            mask_arr[mask_idxs] = np.array(cmap(class_idx)[:3]) * 255

        mask_pil = PIL.Image.fromarray(mask_arr)
        img = PIL.Image.blend(img, mask_pil, alpha=alpha)

    return np.array(img)


def draw_keypoints(
    img: np.ndarray,
    kps: KeyPoints,
    color: Tuple[int, int, int],
):
    x, y, v = kps.x, kps.y, kps.visible

    # calculate scaling for points and connections
    img_h, img_w, _ = img.shape
    img_area = img_h * img_w
    img = PIL.Image.fromarray(img)
    draw = PIL.ImageDraw.Draw(img)
    dynamic_size = int(0.01867599 * (img_area**0.4422045))
    dynamic_size = max(dynamic_size, 1)

    # draw connections
    if kps.metadata is not None and kps.metadata.connections is not None:
        for connection in kps.metadata.connections:
            if v[connection.p1] > 0 and v[connection.p2] > 0:
                draw.line(
                    xy=(
                        (int(x[connection.p1]), int(y[connection.p1])),
                        (int(x[connection.p2]), int(y[connection.p2])),
                    ),
                    fill=connection.color,
                    width=dynamic_size,
                )

    # draw points
    for x_c, y_c in zip(x[v > 0], y[v > 0]):
        radius = dynamic_size
        x1 = x_c - radius
        x2 = x_c + radius
        y1 = y_c - radius
        y2 = y_c + radius
        draw.ellipse(
            xy=[x1, y1, x2, y2],
            # xy=[x1, x0, y1, y0],
            fill=as_rgb_tuple(color),
            outline=None,
            width=0,
        )

    return np.array(img)
