import gradio as gr
import numpy as np
import torch
import cv2
from PIL import Image
import os
from transformers import AutoModelForMaskGeneration, AutoProcessor, pipeline
from dataclasses import dataclass
from typing import List, Dict, Any, Optional
from tqdm import tqdm

@dataclass
class BoundingBox:
    xmin: int
    ymin: int
    ymax: int
    xmax: int

    @property
    def xyxy(self) -> List[float]:
        return [self.xmin, self.ymin, self.xmax, self.ymax]

@dataclass
class DetectionResult:
    score: float
    label: str
    box: BoundingBox
    mask: Optional[np.array] = None

    @classmethod
    def from_dict(cls, detection_dict: Dict) -> 'DetectionResult':
        return cls(
            score=detection_dict['score'],
            label=detection_dict['label'],
            box=BoundingBox(
                xmin=detection_dict['box']['xmin'],
                ymin=detection_dict['box']['ymin'],
                xmax=detection_dict['box']['xmax'],
                ymax=detection_dict['box']['ymax']
            )
        )

def mask_to_polygon(mask: np.ndarray) -> List[List[int]]:
    mask = (mask > 0).astype(np.uint8)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return []
    largest_contour = max(contours, key=cv2.contourArea)
    epsilon = 0.001 * cv2.arcLength(largest_contour, True)
    approx_contour = cv2.approxPolyDP(largest_contour, epsilon, True)
    return approx_contour.reshape(-1, 2).tolist()

def polygon_to_mask(polygon: List, image_shape: tuple) -> np.ndarray:
    if not polygon:
        return np.zeros(image_shape, dtype=np.uint8)
    mask = np.zeros(image_shape, dtype=np.uint8)
    pts = np.array(polygon, dtype=np.int32)
    cv2.fillPoly(mask, [pts], color=(1,))
    return mask

def get_boxes(results: List[DetectionResult]) -> List[List[List[float]]]:
    boxes = []
    for result in results:
        xyxy = result.box.xyxy
        boxes.append(xyxy)
    return [boxes]

def refine_masks(masks: torch.BoolTensor, polygon_refinement: bool = False) -> List[np.ndarray]:
    masks = masks.cpu().float()
    masks = masks.permute(0, 2, 3, 1)
    masks = masks.mean(axis=-1)
    masks = (masks > 0.1).float()
    masks = masks.numpy()
    
    masks = [cv2.GaussianBlur(mask, (3,3), 0) for mask in masks]
    masks = [(mask > 0.2).astype(np.uint8) for mask in masks]

    if polygon_refinement:
        refined_masks = []
        for mask in masks:
            if np.sum(mask) > 100:
                shape = mask.shape
                polygon = mask_to_polygon(mask)
                if polygon:
                    mask = polygon_to_mask(polygon, shape)
                    # Dilate mask slightly to ensure coverage
                    kernel = np.ones((3,3), np.uint8)
                    mask = cv2.dilate(mask, kernel, iterations=1)
            refined_masks.append(mask)
        masks = refined_masks

    return masks

def detect(
    image: Image.Image,
    labels: List[str],
    threshold: float = 0.3,
    detector_id: str = "IDEA-Research/grounding-dino-base"
) -> List[DetectionResult]:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    object_detector = pipeline(model=detector_id, task="zero-shot-object-detection", device=device)
    
    labels = [label if label.endswith(".") else label+"." for label in labels]
    results = object_detector(image, candidate_labels=labels, threshold=threshold)
    results = [DetectionResult.from_dict(result) for result in results]
    
    return results

def segment(
    image: Image.Image,
    detection_results: List[DetectionResult],
    polygon_refinement: bool = False,
    segmenter_id: str = "facebook/sam-vit-base"
) -> List[DetectionResult]:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to(device)
    processor = AutoProcessor.from_pretrained(segmenter_id)

    boxes = get_boxes(detection_results)
    inputs = processor(images=image, input_boxes=boxes, return_tensors="pt").to(device)

    outputs = segmentator(**inputs)
    masks = processor.post_process_masks(
        masks=outputs.pred_masks,
        original_sizes=inputs.original_sizes,
        reshaped_input_sizes=inputs.reshaped_input_sizes
    )[0]

    masks = refine_masks(masks, polygon_refinement)

    for detection_result, mask in zip(detection_results, masks):
        detection_result.mask = mask

    return detection_results

def annotate(image: Image.Image, detection_results: List[DetectionResult]) -> np.ndarray:
    image_cv2 = np.array(image)
    image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)

    for detection in detection_results:
        label = detection.label
        score = detection.score
        box = detection.box
        mask = detection.mask

        color = np.random.randint(0, 256, size=3)

        if mask is not None:
            overlay = image_cv2.copy()
            mask_display = (mask * 255).astype(np.uint8)
            overlay[mask == 1] = color
            image_cv2 = cv2.addWeighted(image_cv2, 0.7, overlay, 0.3, 0)

            contours, _ = cv2.findContours(mask_display, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cv2.drawContours(image_cv2, contours, -1, color.tolist(), 2)

        cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color.tolist(), 2)
        cv2.putText(
            image_cv2,
            f'{label}: {score:.2f}',
            (box.xmin, box.ymin - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            color.tolist(),
            2
        )

    return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)

def process_image(input_image, text_prompt, box_threshold=0.3, save_mask=False):
    try:
        if isinstance(input_image, np.ndarray):
            input_image = Image.fromarray(input_image)
        
        labels = [text_prompt]
        detections = detect(input_image, labels, box_threshold)
        detections = segment(
            input_image,
            detections,
            polygon_refinement=True,
            segmenter_id="facebook/sam-vit-base"
        )

        output_image = annotate(input_image, detections)

        os.makedirs("outputs", exist_ok=True)
        os.makedirs("generated_masks", exist_ok=True)

        mask_image = None
        mask_coordinates = []

        if save_mask and detections and detections[0].mask is not None:
            mask = detections[0].mask
            
            # Create full-size mask
            mask = (mask * 255).astype(np.uint8)
            
            # Apply morphological operations to improve mask
            kernel = np.ones((3,3), np.uint8)
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
            
            mask_image = Image.fromarray(mask, mode="L")

            # Get mask coordinates
            mask_coords = np.where(mask > 0)
            if mask_coords[0].size > 0:
                mask_coordinates = [(int(y), int(x)) for y, x in zip(mask_coords[0], mask_coords[1])]

            if save_mask:
                mask_filename = f"mask_{text_prompt}_{box_threshold}.png"
                mask_image.save(os.path.join("generated_masks", mask_filename))

        output_filename = f"output_{text_prompt}_{box_threshold}.png"
        Image.fromarray(output_image).save(os.path.join("outputs", output_filename))
        
        return output_image, mask_image, mask_coordinates

    except Exception as e:
        print(f"Error during processing: {str(e)}")
        import traceback
        traceback.print_exc()
        return input_image, None, []

def batch_process_images(
        input_folder,
        output_segmented_folder,
        output_masks_folder,
        text_prompt,
        box_threshold=0.3,
):
    os.makedirs(output_segmented_folder, exist_ok=True)
    os.makedirs(output_masks_folder, exist_ok=True)

    image_files = [
        f
        for f in os.listdir(input_folder)
        if os.path.isfile(os.path.join(input_folder, f))
           and f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.webp'))
    ]

    for image_file in tqdm(image_files, desc="Processing images"):
        try:
            input_image_path = os.path.join(input_folder, image_file)
            input_image = Image.open(input_image_path).convert("RGB")
            
            temp_image_path = os.path.join(input_folder, "temp_" + os.path.splitext(image_file)[0] + ".png")
            input_image.save(temp_image_path, format="PNG")

            segmented_image, mask_image, _ = process_image(
                Image.open(temp_image_path), 
                text_prompt, 
                box_threshold, 
                save_mask=True
            )

            segmented_image_path = os.path.join(
                output_segmented_folder, os.path.splitext(image_file)[0] + ".png"
            )
            Image.fromarray(segmented_image).save(segmented_image_path)

            if mask_image:
                mask_image_path = os.path.join(
                    output_masks_folder, os.path.splitext(image_file)[0] + ".png"
                )
                mask_image.save(mask_image_path)
                
            os.remove(temp_image_path)

        except Exception as e:
            print(f"Error processing {image_file}: {e}")

def create_gradio_interface():
    with gr.Blocks() as interface:
        with gr.Tab("Single Image Processing"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type="numpy", label="Input Image")
                    text_prompt = gr.Textbox(
                        label="Text Prompt",
                        value="eyes",
                        placeholder="Enter object to segment (e.g., 'eyes', 'face')",
                    )
                    box_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.3,
                        label="Detection Threshold",
                    )
                    run_button = gr.Button("Process Image")
                with gr.Column():
                    output_image = gr.Image(type="numpy", label="Output Image")
                    output_mask = gr.Image(type="numpy", label="Output Mask")

            def process_and_return_visible(image, prompt, threshold, save_mask):
                output_img, mask_img, _ = process_image(image, prompt, threshold, save_mask)
                return output_img, mask_img

            run_button.click(
                fn=process_and_return_visible,
                inputs=[input_image, text_prompt, box_threshold, gr.Checkbox(True, label="Save Mask")],
                outputs=[output_image, output_mask],
            )

        with gr.Tab("Batch Processing"):
            with gr.Row():
                with gr.Column():
                    input_folder = gr.Textbox(label="Input Folder", placeholder="Enter path to input folder")
                    output_segmented_folder = gr.Textbox(label="Output Segmented Images Folder", placeholder="Enter path to output segmented images folder")
                    output_masks_folder = gr.Textbox(label="Output Masks Folder", placeholder="Enter path to output masks folder")
                    batch_text_prompt = gr.Textbox(
                        label="Text Prompt",
                        value="eyes",
                        placeholder="Enter object to segment (e.g., 'eyes', 'face')",
                    )
                    batch_box_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.3,
                        label="Detection Threshold",
                    )
                    batch_run_button = gr.Button("Batch Process Images")
                with gr.Column():
                    batch_status = gr.Textbox(label="Status", interactive=False)
            
            batch_run_button.click(
                fn=lambda *args: batch_process_images(*args) or "Batch processing completed",
                inputs=[input_folder, output_segmented_folder, output_masks_folder, batch_text_prompt, batch_box_threshold],
                outputs=batch_status
            )

    return interface

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(share=False, inbrowser=True)