import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
import pytesseract
from ultralytics import YOLO
from sklearn.metrics.pairwise import cosine_similarity

# Load CLIP model and processor
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Load YOLOv8 model
yolo = YOLO("yolov8x.pt")  # change to yolov8n.pt for faster inference if needed

# Set Tesseract OCR path if needed
pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"

def extract_clip_features(image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        features = model.get_image_features(**inputs)
    return features / features.norm(p=2, dim=-1, keepdim=True)

def detect_text(image):
    return pytesseract.image_to_string(image)

def detect_objects(image):
    results = yolo(image)
    boxes = results[0].boxes.xywh.cpu().numpy()  # (x, y, w, h)
    return boxes

def compare_images(img1, img2):
    img1 = img1.convert("RGB")
    img2 = img2.convert("RGB")

    # Extract features
    feat1 = extract_clip_features(img1)
    feat2 = extract_clip_features(img2)

    # Cosine similarity
    sim_score = cosine_similarity(feat1, feat2)[0][0]

    # OCR text similarity
    text1 = detect_text(img1)
    text2 = detect_text(img2)

    vec1 = np.array([ord(c) for c in text1[:100]] + [0]*100)[:100]
    vec2 = np.array([ord(c) for c in text2[:100]] + [0]*100)[:100]
    text_sim = cosine_similarity([vec1], [vec2])[0][0]

    # Object detection
    obj1 = detect_objects(np.array(img1))
    obj2 = detect_objects(np.array(img2))
    shape_diff = abs(len(obj1) - len(obj2)) / max(len(obj1), 1)

    # Geometric center distance (avg heuristic)
    dist_penalty = 0
    for i in range(min(len(obj1), len(obj2))):
        dist_penalty += np.linalg.norm(obj1[i][:2] - obj2[i][:2])
    dist_penalty /= max(len(obj1), 1)

    # Final similarity score
    final_score = (0.5 * sim_score) + (0.3 * text_sim) + (0.2 * (1 - shape_diff))
    rating = round(final_score * 5, 2)
    rating_clamped = min(5.0, max(0.0, rating))

    return {
        "Similarity Score (%)": f"{final_score*100:.2f}%",
        "Rating (0–5)": f"{rating_clamped:.1f} ⭐",
        "Text in Image 1": text1.strip()[:200],
        "Text in Image 2": text2.strip()[:200],
        "Detected Objects (img1, img2)": f"{len(obj1)} vs {len(obj2)}"
    }

def gradio_ui(img1, img2):
    result = compare_images(img1, img2)
    return (
        result["Similarity Score (%)"],
        result["Rating (0–5)"],
        result["Text in Image 1"],
        result["Text in Image 2"],
        result["Detected Objects (img1, img2)"]
    )

# Gradio Interface
demo = gr.Interface(
    fn=gradio_ui,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Image(type="pil", label="Output Image")
    ],
    outputs=[
        gr.Text(label="Similarity Score (%)"),
        gr.Text(label="Predicted Rating (0-5 Stars)"),
        gr.Textbox(label="Extracted Text from Input"),
        gr.Textbox(label="Extracted Text from Output"),
        gr.Text(label="Part/Object Count Comparison")
    ],
    title="🛠️ CAD Image Comparison AI",
    description="Upload two CAD/Technical images and get a full feature-based similarity score including shapes, geometry, dimensions, and text."
)

demo.launch()