from typing import Optional

import cv2
import numpy as np
import supervision as sv
from PIL import Image
from transformers import SamImageProcessor, SamModel, SamProcessor, pipeline

from maestro.postprocessing.mask import masks_to_marks


class SegmentAnythingMarkGenerator:
    """A class for performing image segmentation using a specified model.

    Parameters:
        device (str): The device to run the model on (e.g., 'cpu', 'cuda').
        model_name (str): The name of the model to be loaded. Defaults to
            'facebook/sam-vit-huge'.
    """

    def __init__(self, device: str = "cpu", model_name: str = "facebook/sam-vit-huge") -> None:
        self.model = SamModel.from_pretrained(model_name).to(device)
        self.processor = SamProcessor.from_pretrained(model_name)
        self.image_processor = SamImageProcessor.from_pretrained(model_name)
        self.device = device
        self.pipeline = pipeline(
            task="mask-generation", model=self.model, image_processor=self.image_processor, device=self.device
        )

    def generate(self, image: np.ndarray, mask: Optional[np.ndarray] = None) -> sv.Detections:
        """Generate image segmentation marks.

        Parameters:
            image (np.ndarray): The image to be marked in BGR format.
            mask: (Optional[np.ndarray]): The mask to be used as a guide for
                segmentation.

        Returns:
            sv.Detections: An object containing the segmentation masks and their
                corresponding bounding box coordinates.
        """
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        if mask is None:
            outputs = self.pipeline(image, points_per_batch=64)
            masks = np.array(outputs["masks"])
            return masks_to_marks(masks=masks)
        else:
            inputs = self.processor(image, return_tensors="pt").to(self.device)
            image_embeddings = self.model.get_image_embeddings(inputs.pixel_values)
            masks = []
            for polygon in sv.mask_to_polygons(mask.astype(bool)):
                indexes = np.random.default_rng().choice(a=polygon.shape[0], size=5, replace=True)
                input_points = polygon[indexes]
                inputs = self.processor(images=image, input_points=[[input_points]], return_tensors="pt").to(
                    self.device
                )
                del inputs["pixel_values"]
                outputs = self.model(image_embeddings=image_embeddings, **inputs)
                mask = self.processor.image_processor.post_process_masks(
                    masks=outputs.pred_masks.cpu().detach(),
                    original_sizes=inputs["original_sizes"].cpu().detach(),
                    reshaped_input_sizes=inputs["reshaped_input_sizes"].cpu().detach(),
                )[0][0][0].numpy()
                masks.append(mask)
            masks = np.array(masks)
            return masks_to_marks(masks=masks)
