import gradio as gr
import numpy as np
from PIL import Image, ImageDraw


def get_valid_mask(mask: np.ndarray):
    """Convert mask from gr.Image(0 to 255, RGBA) to binary mask.
    """
    if mask.ndim == 3:
        mask_pil = Image.fromarray(mask).convert('L')
        mask = np.array(mask_pil)
    if mask.max() == 255:
        mask = mask / 255
    return mask


def draw_points_on_image(image, points, intr, extr, z, radius_scale=0.006):
    overlay_rgba = Image.new("RGBA", image.size, 0)
    overlay_draw = ImageDraw.Draw(overlay_rgba)
    for point_key, point in points.items():

        t_color = (255, 100, 100)
        o_color = (255, 50, 50)

        rad_draw = int(image.size[0] * radius_scale) + 2

        p_start = point["start"]
        p_target = point["target"]

        if p_start is not None and p_target is not None:
            p_draw = int(p_start[0]), int(p_start[1])
            t_draw = int(p_target[0]), int(p_target[1])

            pt = (p_target[0] - p_start[0], p_target[1] - p_start[1])
            pt_norm = np.linalg.norm(pt)
            pt_unit = (pt[0] / pt_norm, pt[1] / pt_norm)
            pt_tang = (pt_unit[1], -pt_unit[0])
            tt1 = (t_draw[0] + pt_tang[0] * 0.1 * pt_norm - pt_unit[0] * 0.1 * pt_norm,
                   t_draw[1] + pt_tang[1] * 0.1 * pt_norm - pt_unit[1] * 0.1 * pt_norm)
            tt2 = (t_draw[0] - pt_tang[0] * 0.1 * pt_norm - pt_unit[0] * 0.1 * pt_norm,
                   t_draw[1] - pt_tang[1] * 0.1 * pt_norm - pt_unit[1] * 0.1 * pt_norm)
            tt1_draw = int(tt1[0]), int(tt1[1])
            tt2_draw = int(tt2[0]), int(tt2[1])

            overlay_draw.line(
                (p_draw[0], p_draw[1], t_draw[0], t_draw[1]),
                fill=o_color,
                width=4,
            )
            
            overlay_draw.line(
                (t_draw[0], t_draw[1], tt1_draw[0], tt1_draw[1]),
                fill=o_color,
                width=4,
            )

            overlay_draw.line(
                (t_draw[0], t_draw[1], tt2_draw[0], tt2_draw[1]),
                fill=o_color,
                width=4,
            )

        if p_start is not None:
            p_draw = int(p_start[0]), int(p_start[1])
            overlay_draw.ellipse(
                (
                    p_draw[0] - rad_draw,
                    p_draw[1] - rad_draw,
                    p_draw[0] + rad_draw,
                    p_draw[1] + rad_draw,
                ),
                fill=t_color,
                outline=o_color,
                width=2,
            )

        if p_target is not None:
            assert p_start is not None

    return Image.alpha_composite(image.convert("RGBA"),
                                 overlay_rgba).convert("RGB")


def draw_raw_points_on_image(image,
                             points,
                             # curr_point=None,
                             # highlight_all=True,
                             radius_scale=0.002):
    overlay_rgba = Image.new("RGBA", image.size, 0)
    overlay_draw = ImageDraw.Draw(overlay_rgba)
    for p in range(points.shape[0]):
        point = points[p]
        t_color = (150, 150, 255)
        o_color = (50, 50, 255)

        rad_draw = int(image.size[0] * radius_scale)

        t_draw = int(point[0]), int(point[1])
        overlay_draw.ellipse(
            (
                t_draw[0] - rad_draw,
                t_draw[1] - rad_draw,
                t_draw[0] + rad_draw,
                t_draw[1] + rad_draw,
            ),
            fill=t_color,
            outline=o_color,
        )
    return Image.alpha_composite(image.convert("RGBA"),
                                 overlay_rgba).convert("RGB")


def draw_mask_on_image(image, mask):
    im_mask = np.uint8(mask * 255)
    im_mask_rgba = np.concatenate(
        (
            np.tile(im_mask[..., None], [1, 1, 3]),
            45 * np.ones(
                (im_mask.shape[0], im_mask.shape[1], 1), dtype=np.uint8),
        ),
        axis=-1,
    )
    im_mask_rgba = Image.fromarray(im_mask_rgba).convert("RGBA")

    return Image.alpha_composite(image.convert("RGBA"),
                                 im_mask_rgba).convert("RGB")


def on_change_single_global_state(keys,
                                  value,
                                  global_state,
                                  map_transform=None):
    if map_transform is not None:
        value = map_transform(value)

    curr_state = global_state
    if isinstance(keys, str):
        last_key = keys

    else:
        for k in keys[:-1]:
            curr_state = curr_state[k]

        last_key = keys[-1]

    curr_state[last_key] = value
    return global_state


def get_latest_points_pair(points_dict):
    if not points_dict:
        return None
    point_idx = list(points_dict.keys())
    latest_point_idx = max(point_idx)
    return latest_point_idx