import cv2

from src.data import load_data
from src.img_utils import (
    apply_color_overlay,
    apply_composite_overlay,
    apply_gradient_heatmap_overlay,
    apply_spotlight_heatmap,
    blur_background,
    desaturate_background,
    draw_border,
)


def main():
    print("Starting Abduction Demo")
    labeled_images = load_data()
    image = labeled_images[0].image
    mask = labeled_images[0].create_mask([labeled_images[0].labels[0]])

    overlay_result = apply_color_overlay(image, mask, color=(255, 100, 0), alpha=0.4)
    cv2.imwrite(".tmp-data/highlight_overlay.jpg", overlay_result)

    blur_result = blur_background(image, mask, blur_intensity=(51, 51))
    cv2.imwrite(".tmp-data/highlight_blurred.jpg", blur_result)

    desaturate_result = desaturate_background(image, mask)
    cv2.imwrite(".tmp-data/highlight_desaturated.jpg", desaturate_result)

    gradient_heatmap_result = apply_gradient_heatmap_overlay(image, mask, colormap=cv2.COLORMAP_JET, alpha=0.6)
    cv2.imwrite(".tmp-data/highlight_gradient_heatmap.jpg", gradient_heatmap_result)

    spotlight_result = apply_spotlight_heatmap(image, mask, darkness_factor=0.2)
    cv2.imwrite(".tmp-data/highlight_spotlight_heatmap.jpg", spotlight_result)

    composite_result = apply_composite_overlay(
        image, mask, colormap=cv2.COLORMAP_JET, foreground_alpha=0.6, background_alpha=0.5
    )
    cv2.imwrite(".tmp-data/highlight_composite.jpg", composite_result)

    bordered_image = draw_border(
        image,
        mask,
        color=(50, 255, 50),  # A bright green color
        thickness=8,
    )
    cv2.imwrite(".tmp-data/highlight_border.jpg", bordered_image)


if __name__ == "__main__":
    main()