# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved

import colorsys
from dataclasses import dataclass
from typing import List, Tuple

import cv2
import matplotlib as mpl
import matplotlib.colors as mplc
import numpy as np
import pycocotools.mask as mask_utils


def rgb_to_hex(rgb_color):
    """
    Convert a rgb color to hex color.

    Args:
        rgb_color (tuple/list of ints): RGB color in tuple or list format.

    Returns:
        str: Hex color.

    Example:
        ```
        >>> rgb_to_hex((255, 0, 244))
        '#ff00ff'
        ```
    """
    return "#" + "".join([hex(c)[2:].zfill(2) for c in rgb_color])


# DEFAULT_COLOR_HEX_TO_NAME = {
#     rgb_to_hex((255, 0, 0)): "red",
#     rgb_to_hex((0, 255, 0)): "lime",
#     rgb_to_hex((0, 0, 255)): "blue",
#     rgb_to_hex((255, 255, 0)): "yellow",
#     rgb_to_hex((255, 0, 255)): "fuchsia",
#     rgb_to_hex((0, 255, 255)): "aqua",
#     rgb_to_hex((255, 165, 0)): "orange",
#     rgb_to_hex((128, 0, 128)): "purple",
#     rgb_to_hex((255, 215, 0)): "gold",
# }

# Assuming rgb_to_hex is a function that converts an (R, G, B) tuple to a hex string.
# For example: def rgb_to_hex(rgb): return '#%02x%02x%02x' % rgb

DEFAULT_COLOR_HEX_TO_NAME = {
    # The top 20 approved colors
    rgb_to_hex((255, 255, 0)): "yellow",
    rgb_to_hex((0, 255, 0)): "lime",
    rgb_to_hex((0, 255, 255)): "cyan",
    rgb_to_hex((255, 0, 255)): "magenta",
    rgb_to_hex((255, 0, 0)): "red",
    rgb_to_hex((255, 127, 0)): "orange",
    rgb_to_hex((127, 255, 0)): "chartreuse",
    rgb_to_hex((0, 255, 127)): "spring green",
    rgb_to_hex((255, 0, 127)): "rose",
    rgb_to_hex((127, 0, 255)): "violet",
    rgb_to_hex((192, 255, 0)): "electric lime",
    rgb_to_hex((255, 192, 0)): "vivid orange",
    rgb_to_hex((0, 255, 192)): "turquoise",
    rgb_to_hex((192, 0, 255)): "bright violet",
    rgb_to_hex((255, 0, 192)): "bright pink",
    rgb_to_hex((255, 64, 0)): "fiery orange",
    rgb_to_hex((64, 255, 0)): "bright chartreuse",
    rgb_to_hex((0, 255, 64)): "malachite",
    rgb_to_hex((64, 0, 255)): "deep violet",
    rgb_to_hex((255, 0, 64)): "hot pink",
}


DEFAULT_COLOR_PALETTE = list(DEFAULT_COLOR_HEX_TO_NAME.keys())


def _validate_color_hex(color_hex: str):
    color_hex = color_hex.lstrip("#")
    if not all(c in "0123456789abcdefABCDEF" for c in color_hex):
        raise ValueError("Invalid characters in color hash")
    if len(color_hex) not in (3, 6):
        raise ValueError("Invalid length of color hash")


# copied from https://github.com/roboflow/supervision/blob/c8f557af0c61b5c03392bad2cc36c8835598b1e1/supervision/draw/color.py
@dataclass
class Color:
    """
    Represents a color in RGB format.

    Attributes:
        r (int): Red channel.
        g (int): Green channel.
        b (int): Blue channel.
    """

    r: int
    g: int
    b: int

    @classmethod
    def from_hex(cls, color_hex: str):
        """
        Create a Color instance from a hex string.

        Args:
            color_hex (str): Hex string of the color.

        Returns:
            Color: Instance representing the color.

        Example:
            ```
            >>> Color.from_hex('#ff00ff')
            Color(r=255, g=0, b=255)
            ```
        """
        _validate_color_hex(color_hex)
        color_hex = color_hex.lstrip("#")
        if len(color_hex) == 3:
            color_hex = "".join(c * 2 for c in color_hex)
        r, g, b = (int(color_hex[i : i + 2], 16) for i in range(0, 6, 2))
        return cls(r, g, b)

    @classmethod
    def to_hex(cls, color):
        """
        Convert a Color instance to a hex string.

        Args:
            color (Color): Color instance of color.

        Returns:
            Color: a hex string.
        """
        return rgb_to_hex((color.r, color.g, color.b))

    def as_rgb(self) -> Tuple[int, int, int]:
        """
        Returns the color as an RGB tuple.

        Returns:
            Tuple[int, int, int]: RGB tuple.

        Example:
            ```
            >>> color.as_rgb()
            (255, 0, 255)
            ```
        """
        return self.r, self.g, self.b

    def as_bgr(self) -> Tuple[int, int, int]:
        """
        Returns the color as a BGR tuple.

        Returns:
            Tuple[int, int, int]: BGR tuple.

        Example:
            ```
            >>> color.as_bgr()
            (255, 0, 255)
            ```
        """
        return self.b, self.g, self.r

    @classmethod
    def white(cls):
        return Color.from_hex(color_hex="#ffffff")

    @classmethod
    def black(cls):
        return Color.from_hex(color_hex="#000000")

    @classmethod
    def red(cls):
        return Color.from_hex(color_hex="#ff0000")

    @classmethod
    def green(cls):
        return Color.from_hex(color_hex="#00ff00")

    @classmethod
    def blue(cls):
        return Color.from_hex(color_hex="#0000ff")


@dataclass
class ColorPalette:
    colors: List[Color]

    @classmethod
    def default(cls):
        """
        Returns a default color palette.

        Returns:
            ColorPalette: A ColorPalette instance with default colors.

        Example:
            ```
            >>> ColorPalette.default()
            ColorPalette(colors=[Color(r=255, g=0, b=0), Color(r=0, g=255, b=0), ...])
            ```
        """
        return ColorPalette.from_hex(color_hex_list=DEFAULT_COLOR_PALETTE)

    @classmethod
    def from_hex(cls, color_hex_list: List[str]):
        """
        Create a ColorPalette instance from a list of hex strings.

        Args:
            color_hex_list (List[str]): List of color hex strings.

        Returns:
            ColorPalette: A ColorPalette instance.

        Example:
            ```
            >>> ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff'])
            ColorPalette(colors=[Color(r=255, g=0, b=0), Color(r=0, g=255, b=0), ...])
            ```
        """
        colors = [Color.from_hex(color_hex) for color_hex in color_hex_list]
        return cls(colors)

    def by_idx(self, idx: int) -> Color:
        """
        Return the color at a given index in the palette.

        Args:
            idx (int): Index of the color in the palette.

        Returns:
            Color: Color at the given index.

        Example:
            ```
            >>> color_palette.by_idx(1)
            Color(r=0, g=255, b=0)
            ```
        """
        if idx < 0:
            raise ValueError("idx argument should not be negative")
        idx = idx % len(self.colors)
        return self.colors[idx]

    def find_farthest_color(self, img_array):
        """
        Return the color that is the farthest from the given color.

        Args:
            img_array (np array): any *x3 np array, 3 is the RGB color channel.

        Returns:
            Color: Farthest color.

        """
        # Reshape the image array for broadcasting
        img_array = img_array.reshape((-1, 3))

        # Convert colors dictionary to a NumPy array
        color_values = np.array([[c.r, c.g, c.b] for c in self.colors])

        # Calculate the Euclidean distance between the colors and each pixel in the image
        # Broadcasting happens here: img_array shape is (num_pixels, 3), color_values shape is (num_colors, 3)
        distances = np.sqrt(
            np.sum((img_array[:, np.newaxis, :] - color_values) ** 2, axis=2)
        )

        # Average the distances for each color
        mean_distances = np.mean(distances, axis=0)

        # return the farthest color
        farthest_idx = np.argmax(mean_distances)
        farthest_color = self.colors[farthest_idx]
        farthest_color_hex = Color.to_hex(farthest_color)
        if farthest_color_hex in DEFAULT_COLOR_HEX_TO_NAME:
            farthest_color_name = DEFAULT_COLOR_HEX_TO_NAME[farthest_color_hex]
        else:
            farthest_color_name = "unknown"

        return farthest_color, farthest_color_name


def draw_box(ax, box_coord, alpha=0.8, edge_color="g", line_style="-", linewidth=2.0):
    x0, y0, width, height = box_coord
    ax.add_patch(
        mpl.patches.Rectangle(
            (x0, y0),
            width,
            height,
            fill=False,
            edgecolor=edge_color,
            linewidth=linewidth,
            alpha=alpha,
            linestyle=line_style,
        )
    )


def draw_text(
    ax,
    text,
    position,
    font_size=None,
    color="g",
    horizontal_alignment="left",
    rotation=0,
):
    if not font_size:
        font_size = mpl.rcParams["font.size"]

    color = np.maximum(list(mplc.to_rgb(color)), 0.2)
    color[np.argmax(color)] = max(0.8, np.max(color))

    x, y = position
    ax.text(
        x,
        y,
        text,
        size=font_size,
        family="sans-serif",
        bbox={"facecolor": "none", "alpha": 0.5, "pad": 0.7, "edgecolor": "none"},
        verticalalignment="top",
        horizontalalignment=horizontal_alignment,
        color=color,
        rotation=rotation,
    )


def draw_mask(
    ax, rle, color, show_holes=True, alpha=0.15, upsample_factor=1.0, rle_upsampled=None
):
    if isinstance(rle, dict):
        mask = mask_utils.decode(rle)
    elif isinstance(rle, np.ndarray):
        mask = rle
    else:
        raise ValueError(f"Unsupported type for rle: {type(rle)}")

    mask_upsampled = None
    if upsample_factor > 1.0 and show_holes:
        assert rle_upsampled is not None
        if isinstance(rle_upsampled, dict):
            mask_upsampled = mask_utils.decode(rle_upsampled)
        elif isinstance(rle_upsampled, np.ndarray):
            mask_upsampled = rle_upsampled
        else:
            raise ValueError(f"Unsupported type for rle: {type(rle)}")

    if show_holes:
        if mask_upsampled is None:
            mask_upsampled = mask
        h, w = mask_upsampled.shape
        mask_img = np.zeros((h, w, 4))
        mask_img[:, :, :-1] = color[np.newaxis, np.newaxis, :]
        mask_img[:, :, -1] = mask_upsampled * alpha
        ax.imshow(mask_img)

    *_, contours, _ = cv2.findContours(
        mask.astype(np.uint8).copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
    )
    upsampled_contours = [(cont + 0.5) * upsample_factor - 0.5 for cont in contours]
    facecolor = (0, 0, 0, 0) if show_holes else color
    if alpha > 0.8:
        edge_color = _change_color_brightness(color, brightness_factor=-0.7)
    else:
        edge_color = color
    for cont in upsampled_contours:
        polygon = mpl.patches.Polygon(
            [el[0] for el in cont],
            edgecolor=edge_color,
            linewidth=2.0,
            facecolor=facecolor,
        )
        ax.add_patch(polygon)


def _change_color_brightness(color, brightness_factor):
    """
    Depending on the brightness_factor, gives a lighter or darker color i.e. a color with
    less or more saturation than the original color.

    Args:
        color: color of the polygon. Refer to `matplotlib.colors` for a full list of
            formats that are accepted.
        brightness_factor (float): a value in [-1.0, 1.0] range. A lightness factor of
            0 will correspond to no change, a factor in [-1.0, 0) range will result in
            a darker color and a factor in (0, 1.0] range will result in a lighter color.

    Returns:
        modified_color (tuple[double]): a tuple containing the RGB values of the
            modified color. Each value in the tuple is in the [0.0, 1.0] range.
    """
    assert brightness_factor >= -1.0 and brightness_factor <= 1.0
    color = mplc.to_rgb(color)
    polygon_color = colorsys.rgb_to_hls(*mplc.to_rgb(color))
    modified_lightness = polygon_color[1] + (brightness_factor * polygon_color[1])
    modified_lightness = 0.0 if modified_lightness < 0.0 else modified_lightness
    modified_lightness = 1.0 if modified_lightness > 1.0 else modified_lightness
    modified_color = colorsys.hls_to_rgb(
        polygon_color[0], modified_lightness, polygon_color[2]
    )
    return modified_color