# source: huggingface: fashn-ai/sapiens-body-part-segmentation
import colorsys
import matplotlib.colors as mcolors
import numpy as np
from PIL import Image

def get_palette(num_cls):
    palette = [0] * (256 * 3)
    palette[0:3] = [0, 0, 0]

    for j in range(1, num_cls):
        hue = (j - 1) / (num_cls - 1)
        saturation = 1.0
        value = 1.0 if j % 2 == 0 else 0.5
        rgb = colorsys.hsv_to_rgb(hue, saturation, value)
        r, g, b = [int(x * 255) for x in rgb]
        palette[j * 3 : j * 3 + 3] = [r, g, b]

    return palette


def create_colormap(palette):
    colormap = np.array(palette).reshape(-1, 3) / 255.0
    return mcolors.ListedColormap(colormap)


def visualize_mask_with_overlay(img: Image.Image, mask: Image.Image, labels_to_ids: dict[str, int], alpha=0.5):
    img_np = np.array(img.convert("RGB"))
    mask_np = np.array(mask)

    num_cls = len(labels_to_ids)
    palette = get_palette(num_cls)
    colormap = create_colormap(palette)

    overlay = np.zeros((*mask_np.shape, 3), dtype=np.uint8)
    for label, idx in labels_to_ids.items():
        if idx != 0:
            overlay[mask_np == idx] = np.array(colormap(idx)[:3]) * 255

    blended = Image.fromarray(np.uint8(img_np * (1 - alpha) + overlay * alpha))

    return blended