import gradio as gr
from ultralytics import YOLO
from PIL import Image
import tempfile
import os
import shutil
import numpy as np
import easyocr
from collections import defaultdict
import cv2

# Load YOLO model
print("[INFO] Loading Custom YOLO Segmentation model...")
model = YOLO('trained_seg_model/seg_v1.pt')
print("[INFO] Custom YOLO Segmentation model loaded successfully.")

# Load EasyOCR model
print("[INFO] Loading EasyOCR model...")
reader = easyocr.Reader(['fr'])
print("[INFO] EasyOCR model initialized.")

# OCR helper function
def run_easyocr_on_yolo_segments(results, original_img: Image.Image):
    print("[INFO] Running EasyOCR on segmented regions...")
    original_img_np = np.array(original_img.convert("RGB"))
    ocr_dict = defaultdict(list)
    class_counts = defaultdict(int)
    cropped_images = []

    # for result in results:
    #     if result.masks is None or result.boxes is None:
    #         continue

    for result_index, result in enumerate(results):
        print(f"[DEBUG] Processing result index: {result_index}")
        if result.masks is None or result.boxes is None:
            print("[WARN] Result has no masks or boxes.")
            continue

        # for seg, class_id in zip(result.masks.xy, result.boxes.cls):
        for seg_index, (seg, class_id) in enumerate(zip(result.masks.xy, result.boxes.cls)):
            class_name = model.names[int(class_id)]
            print(f"[DEBUG] Segment {seg_index}: Class = {class_name}")
            poly = np.array(seg, dtype=np.int32)
            x, y, w, h = cv2.boundingRect(poly)
            crop = original_img_np[y:y+h, x:x+w]

            if crop.size == 0:
                print(f"[WARN] Empty crop for class: {class_name}")
                text = ""
            else:
                ocr_result = reader.readtext(crop, detail=0)
                text = " ".join(ocr_result).strip()
                print(f"[DEBUG] OCR result for {class_name}: {text[:60]}...")

            # Always use a suffix, starting from 1
            class_counts[class_name] += 1
            class_key = f"{class_name}_{class_counts[class_name]}"
            ocr_dict[class_key] = text

            # Convert cropped numpy image to PIL for gallery
            cropped_pil = Image.fromarray(crop)
            cropped_images.append((class_key, cropped_pil))

    print(f"[INFO] OCR completed. Total regions processed: {len(cropped_images)}")
    return cropped_images, dict(ocr_dict)

# Main function for Gradio
def segment_and_ocr(img: Image.Image):
    print("[INFO] Starting segmentation + OCR pipeline...")
    with tempfile.TemporaryDirectory() as tmpdir:
        input_path = os.path.join(tmpdir, "input.jpg")
        img.save(input_path)
        print(f"[INFO] Saved input image to: {input_path}")

        print("[INFO] Running YOLO prediction...")
        results = model.predict(
            source=input_path,
            save=True,
            save_txt=True,
            project=tmpdir,
            name="predict",
            exist_ok=True
        )
        print("[INFO] YOLO prediction completed.")
        output_img_path = os.path.join(tmpdir, "predict", "input.jpg")
        label_txt_path = os.path.join(tmpdir, "predict", "labels", "input.txt")

        # segmented_img = Image.open(output_img_path) if os.path.exists(output_img_path) else None
        # label_data = open(label_txt_path).read() if os.path.exists(label_txt_path) else "No labels generated."


        segmented_img = None
        if os.path.exists(output_img_path):
            segmented_img = Image.open(output_img_path)
            print(f"[INFO] Segmented image found: {output_img_path}")
        else:
            print("[ERROR] Segmented image not found!")

        if os.path.exists(label_txt_path):
            with open(label_txt_path) as f:
                label_data = f.read()
            print(f"[INFO] Label file found with length: {len(label_data)} chars")
        else:
            label_data = "No labels generated."
            print("[WARN] Label file not found.")

        print("[INFO] Extracting OCR results from segments...")
        cropped_images_with_labels, ocr_dict = run_easyocr_on_yolo_segments(results, img)
        

        # Format gallery for Gradio [(image, label)]
        gallery_output = [(image, label) for label, image in cropped_images_with_labels]
        print(f"[INFO] Returning {len(gallery_output)} cropped images and {len(ocr_dict)} OCR entries.")


        return segmented_img, label_data, gallery_output, ocr_dict


# Sample images for "Examples" section
examples = [
    ["sample_images/seg_v1.jpg"],
    ["sample_images/seg_v2.jpg"],
    ["sample_images/seg_v3.jpg"],
    ["sample_images/seg_v4.jpg"]
]

# Gradio UI
print("[INFO] Launching Gradio interface...")
gr.Interface(
    fn=segment_and_ocr,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=[
        gr.Image(type="pil", label="Segmented Image"),
        gr.Textbox(label="YOLO Labels (.txt Output)"),
        gr.Gallery(label="Segmented Crops (OCR Regions)"),
        gr.JSON(label="OCR Output by Class (with suffix)")
    ],
    examples=examples,
    cache_examples=False,     # <- this disables automatic execution
    title="AI-Powered YOLO Segmentation + OCR with EasyOCR",
    description="Upload or choose an image to segment using YOLO and extract text using EasyOCR from each region."
).launch()