# app.py
import cv2
import numpy as np
import pytesseract
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
import torch
import os
import uuid
import gradio as gr

# Try importing pptx; give an error message if unavailable
try:
    from pptx import Presentation
    from pptx.util import Inches
except ImportError:
    raise ImportError("Missing 'python-pptx'. Please install it using 'pip install python-pptx'")

# Load CLIP
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Optional: Update Tesseract path if needed
# pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe"

def match_shapes(img1, img2):
    def preprocess(img):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        return cv2.Canny(gray, 100, 200)

    c1 = preprocess(img1)
    c2 = preprocess(img2)

    contours1, _ = cv2.findContours(c1, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours2, _ = cv2.findContours(c2, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    matches = 0
    for cnt1 in contours1:
        for cnt2 in contours2:
            score = cv2.matchShapes(cnt1, cnt2, 1, 0.0)
            if score < 0.2:
                matches += 1
                break

    score = matches / max(len(contours1), 1)
    return round(score * 100, 2), contours1, contours2

def extract_text(img):
    return pytesseract.image_to_string(img).strip()

def compare_text(text1, text2):
    t1_words = set(text1.lower().split())
    t2_words = set(text2.lower().split())
    if not t1_words or not t2_words:
        return 0
    return round(len(t1_words & t2_words) / len(t1_words | t2_words) * 100, 2)

def semantic_score(img1, img2):
    inputs = clip_processor(images=[img1, img2], return_tensors="pt")
    outputs = clip_model.get_image_features(**inputs)
    sim = torch.nn.functional.cosine_similarity(outputs[0][0], outputs[0][1], dim=0)
    return round(float(sim.item()) * 100, 2)

def draw_comparison(img1, img2, contours1, contours2):
    out1 = img1.copy()
    out2 = img2.copy()

    for cnt in contours1:
        cv2.drawContours(out1, [cnt], -1, (0, 0, 255), 2)
    for cnt in contours2:
        cv2.drawContours(out2, [cnt], -1, (0, 255, 0), 2)

    combined = np.hstack((out1, out2))
    return combined

def generate_ppt(source_img_path, comp_img_path, comparison_img, shape_score, text_score, semantic_score):
    prs = Presentation()
    slide = prs.slides.add_slide(prs.slide_layouts[5])
    slide.shapes.title.text = "Image Comparison Report"

    slide1 = prs.slides.add_slide(prs.slide_layouts[5])
    slide1.shapes.title.text = "Input Images"
    slide1.shapes.add_picture(source_img_path, Inches(0.5), Inches(1), width=Inches(4))
    slide1.shapes.add_picture(comp_img_path, Inches(5), Inches(1), width=Inches(4))

    slide2 = prs.slides.add_slide(prs.slide_layouts[5])
    slide2.shapes.title.text = "Comparison Scores"
    scores = f"""Shape Score: {shape_score}% (Rating: {int(shape_score // 20)}/5)\nText Score: {text_score}% (Rating: {int(text_score // 20)}/5)\nSemantic Score: {semantic_score}% (Rating: {int(semantic_score // 20)}/5)"""
    tf = slide2.shapes.add_textbox(Inches(1), Inches(1), Inches(8), Inches(3)).text_frame
    tf.text = scores

    slide3 = prs.slides.add_slide(prs.slide_layouts[5])
    slide3.shapes.title.text = "Visual Comparison"
    fname = f"compare_{uuid.uuid4().hex[:6]}.png"
    cv2.imwrite(fname, comparison_img)
    slide3.shapes.add_picture(fname, Inches(1), Inches(1), width=Inches(7))

    output_ppt = f"comparison_report_{uuid.uuid4().hex[:6]}.pptx"
    prs.save(output_ppt)
    return output_ppt

def process_images(img1_np, img2_np):
    cv2.imwrite("source.png", img1_np)
    cv2.imwrite("target.png", img2_np)
    img1 = cv2.imread("source.png")
    img2 = cv2.imread("target.png")

    shape_score, cnt1, cnt2 = match_shapes(img1, img2)
    text_score = compare_text(extract_text(img1), extract_text(img2))
    sem_score = semantic_score(Image.fromarray(img1_np), Image.fromarray(img2_np))

    compare_img = draw_comparison(img1, img2, cnt1, cnt2)
    ppt_path = generate_ppt("source.png", "target.png", compare_img, shape_score, text_score, sem_score)

    return ppt_path

def gradio_interface(img1, img2):
    ppt = process_images(np.array(img1), np.array(img2))
    return ppt

demo = gr.Interface(
    fn=gradio_interface,
    inputs=[gr.Image(type="pil"), gr.Image(type="pil")],
    outputs=gr.File(label="Download PPT Report"),
    title="Smart CAD Image Comparator",
    description="Upload two technical/CAD images to compare shapes, text, dimensions, and get a detailed PPT report."
)

demo.launch()