from typing import List
from spatstatapp.paths import find_label_path
from spatstatapp.tile_training_data import load_bboxes


import cv2
import numpy as np
import supervision as sv
from ultralytics.utils.ops import xywhn2xyxy


def coco_to_detections(label_path, image, id_to_label = {0:"scallop"}):
    bboxes = load_bboxes(label_path)
    class_id = np.array(bboxes)[:,0].astype(int)
    class_name = np.vectorize(id_to_label.get)(class_id)

    bboxes = np.array(bboxes)[:,1:] # remove class_id
    xyxy = xywhn2xyxy(bboxes,
                    w = image.shape[1],
                    h = image.shape[0])
    xyxy = np.clip(xyxy, a_min=0, a_max=max(image.shape))



    detections = sv.Detections(
        xyxy=xyxy,
        # class_id=np.array(len(xyxy)*[0]),
        class_id=class_id,
        data = {"class_name":class_name}

    )
    return detections

def draw_img_bboxes(img_path):
    image = cv2.imread(img_path)

    label_path = find_label_path(img_path)

    detections = coco_to_detections(label_path, image)

    box_annotator = sv.BoxAnnotator()

    annotated_frame = box_annotator.annotate(
        scene=image.copy(),
        detections=detections
    )
    return annotated_frame

def stack_images_grid(images):
    num_images = len(images)
    grid_size = int(np.ceil(np.sqrt(num_images)))

    # Calculate the size of the grid
    img_height, img_width, _ = images[0].shape
    buffer = 2
    grid_height = grid_size * img_height + (grid_size - 1) * buffer
    grid_width = grid_size * img_width + (grid_size - 1) * buffer

    # Create a white canvas
    grid_image = np.ones((grid_height, grid_width, 3), dtype=np.uint8) * 255

    for idx, img in enumerate(images):
        row = idx // grid_size
        col = idx % grid_size
        y = row * (img_height + buffer)
        x = col * (img_width + buffer)
        grid_image[y:y+img_height, x:x+img_width, :] = img

    return grid_image

def show_img_bboxes(img_path):
    annotated_frame = draw_img_bboxes(img_path)
    sv.plot_image(annotated_frame)




def annotate_DetectionDataset(dataset: sv.DetectionDataset, idxs:List[int]|int):
    if isinstance(idxs, int):
        idxs = np.random.choice(len(dataset), size=idxs, replace=False).tolist()

    frames = []
    for i in idxs:
        path, img, detections = dataset[i]
        box_annotator = sv.BoxAnnotator()
        frames.append(box_annotator.annotate(
                scene=img.copy(),
                detections=detections
            ))
    return frames

def show_img_bboxes_grid(img_paths):
    """
    shows a grid of images with bounding boxes from list of paths

    Args:
        img_paths: list of image paths

    Example:
        show_img_bboxes_grid((data_dir/"images/train").glob("*.jpg")[0:16])

    """
    images = [draw_img_bboxes(img_path) for img_path in img_paths]
    annotated_grid = stack_images_grid(images)
    sv.plot_image(annotated_grid)

def show_dataset_bboxes_grid(dataset: sv.DetectionDataset, idxs:List[int]|int):
    """
    shows a grid of images with bounding boxes from the dataset

    Args:
        dataset: sv.DetectionDataset: _description_
        idxs: list of indices to show or int: number of random indices to show
    Example:
        data_dir = Path('dataset_clipped')

        dd = sv.DetectionDataset.from_yolo(
            images_directory_path= data_dir/"images/train",
            annotations_directory_path= data_dir/"labels/train",
            data_yaml_path="data.yaml",
                )

        show_dataset_bboxes_grid(dd, 16)

    """
    images = annotate_DetectionDataset(dataset, idxs)
    annotated_grid = stack_images_grid(images)
    sv.plot_image(annotated_grid)