# This file is used to visualize bounding boxes on an image
from urllib.parse import urlparse
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import requests
from typing import List
from functools import cache

DEFAULTS = {
  'bbox_outline_width': 2,
  'bbox_outline_color': (0, 0, 256, 123), # alpha runs from 0 to 127
  'bbox_fill_color': (256, 0, 0, 50), # alpha runs from 0 to 127
  'label_text_color': "black",
  'label_fill_color': "red",
  'label_text_padding': 0,
  'label_rectangle_left_margin': 0,
  'label_rectangle_top_margin': 0,
  'label_text_size': 12,
}

@cache
def get_font(path_or_url: str = 'https://github.com/googlefonts/roboto/raw/main/src/hinted/Roboto-Regular.ttf', size: int = DEFAULTS['label_text_size']):
  if urlparse(path_or_url).scheme in ["http", "https"]: # Online
      return ImageFont.truetype(requests.get(path_or_url, stream=True).raw, size=size)
  else: # Local
    return ImageFont.truetype(path_or_url, size=size)

def visualize_bboxes_on_image(
    image: Image.Image,
    bboxes: List[List[int]],
    labels: List[str] = None,
    bbox_outline_width = DEFAULTS["bbox_outline_width"],
    bbox_outline_color = DEFAULTS["bbox_outline_color"],
    bbox_fill_color: str | list[tuple | str] = DEFAULTS["bbox_fill_color"],
    label_text_color: str | list[tuple | str] = DEFAULTS["label_text_color"],
    label_fill_color = DEFAULTS["label_fill_color"],
    label_text_padding = DEFAULTS["label_text_padding"],
    label_rectangle_left_margin = DEFAULTS["label_rectangle_left_margin"],
    label_rectangle_top_margin = DEFAULTS['label_rectangle_top_margin'],
    label_text_size = DEFAULTS["label_text_size"],
    convert_to_x0y0x1y1 = None) -> Image.Image:
  '''
  Visualize bounding boxes on an image
  Args:
    image: Image to visualize
    bboxes: List of bounding boxes
    labels: Titles of the bounding boxes
    bbox_outline_width: Width of the bounding box
    bbox_outline_color: Color of the bounding box
    bbox_fill_color: Fill color of the bounding box
    label_text_color: Color of the label text
    label_fill_color: Color of the label rectangle
    label_text_padding: Padding of the label text
    label_rectangle_left_margin: Left padding of the label rectangle
    label_rectangle_top_margin: Top padding of the label rectangle
    label_text_size: Font size of the label text
    convert_to_x0y0x1y1: Function to convert bounding box to x0y0x1y1 format
  Returns:
    Image: Image annotated with bounding boxes
  '''
  image = image.copy().convert("RGB")
  draw = ImageDraw.Draw(image)
  font = get_font(size = label_text_size)
  labels = (labels or []) + np.full(len(bboxes) - len(labels or []), None).tolist()
  bbox_fill_colors = bbox_fill_color if isinstance(bbox_fill_color, list) else [bbox_fill_color] * len(bboxes)
  bbox_outline_colors = bbox_outline_color if isinstance(bbox_outline_color, list) else [bbox_outline_color] * len(bboxes)
  for bbox, label, _bbox_fill_color, _bbox_outline_color in zip(bboxes, labels, bbox_fill_colors, bbox_outline_colors):
    x0, y0, x1, y1 = convert_to_x0y0x1y1(bbox) if convert_to_x0y0x1y1 is not None else bbox

    rectangle_image = Image.new('RGBA', image.size)
    rectangle_image_draw = ImageDraw.Draw(rectangle_image)
    rectangle_image_draw.rectangle(
      xy = [x0, y0, x1, y1],
      fill = _bbox_fill_color, 
      outline = _bbox_outline_color, 
      width = bbox_outline_width)
    image.paste(im = rectangle_image, mask = rectangle_image)

    if label is not None:
      draw_text_on_image(
        draw, 
        [x0, y0], 
        label, 
        label_text_color, 
        label_fill_color, 
        label_text_padding, 
        label_rectangle_left_margin, 
        label_rectangle_top_margin, 
        label_text_size, 
        font)
  return image

def draw_text_on_image(
    image_or_draw: Image.Image | ImageDraw.ImageDraw,
    text_position_xy: List[int],
    label: str,
    label_text_color = DEFAULTS["label_text_color"],
    label_fill_color = DEFAULTS["label_fill_color"],
    label_text_padding = DEFAULTS["label_text_padding"],
    label_rectangle_left_margin = DEFAULTS["label_rectangle_left_margin"],
    label_rectangle_top_margin = DEFAULTS['label_rectangle_top_margin'],
    label_text_size = DEFAULTS["label_text_size"],
    font: ImageFont.FreeTypeFont = None) -> Image.Image:
  is_image = isinstance(image_or_draw, Image.Image)
  image = image_or_draw.copy().convert("RGB") if is_image else None
  font = font or get_font(size = label_text_size)
  x0, y0 = text_position_xy
  text_position = (x0 - label_rectangle_left_margin + label_text_padding, y0 - label_rectangle_top_margin + label_text_padding)
  draw = ImageDraw.Draw(image) if is_image else image_or_draw
  _, _, text_bbox_right, text_bbox_bottom = draw.textbbox(text_position, label, font=font)
  xy = [
    text_position[0] - label_text_padding, 
    text_position[1] - label_text_padding, 
    text_bbox_right + label_text_padding + label_text_padding, 
    text_bbox_bottom + label_text_padding + label_text_padding
  ]
  draw.rectangle(xy, fill = label_fill_color)
  draw.text(text_position, label, font=font, fill=label_text_color)
  return image