app.py

import gradio as gr
import fitz  # PyMuPDF
import tempfile
import os
import base64
from gtts import gTTS
from ultralytics import YOLO
from PIL import Image, ImageDraw
import io
import json
from datetime import datetime
import random
import numpy as np

# ============================================
# 1. LOAD PDF AND YOLO MODEL
# ============================================

PDF_PATH = "2VBMAPP.pdf"

class ABATutor:
    def __init__(self):
        # Set writable directory for Ultralytics
        os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
        
        # Load YOLOv8 model
        try:
            self.yolo_model = YOLO("yolov8n.pt")
            print("✅ YOLOv8 model loaded")
        except Exception as e:
            print(f"⚠️ YOLO error: {e}")
            self.yolo_model = None
        
        # Load PDF
        try:
            self.doc = fitz.open(PDF_PATH)
            self.total_pages = len(self.doc)
            print(f"✅ Loaded PDF with {self.total_pages} pages")
        except Exception as e:
            print(f"⚠️ PDF error: {e}")
            self.doc = None
            self.total_pages = 1
        
        self.current_page = 6
        self.session_data = []
        self.current_target = None
        self.current_detections = []
        self.current_image_pil = None
        
        # Target objects for each page
        self.page_targets = {
            6: ["cat", "dog", "ball", "shoe", "car", "apple", "bird", "fish", "hat", "cup"],
            7: ["cat", "dog", "ball", "shoe", "car", "apple"],
            8: ["bird", "fish", "hat", "cup", "book", "chair"],
            9: ["train", "plane", "boat", "truck", "bus", "bike"],
            10: ["frog", "monkey", "elephant", "lion", "giraffe", "zebra"],
        }
    
    def get_current_image_with_boxes(self):
        """Get current page image with YOLO bounding boxes"""
        if self.doc is None:
            return Image.new('RGB', (800, 600), color='white'), []
        
        page = self.doc.load_page(self.current_page)
        pix = page.get_pixmap(dpi=120)
        img_data = pix.tobytes("png")
        pil_img = Image.open(io.BytesIO(img_data))
        self.current_image_pil = pil_img.copy()
        
        # Run YOLO
        detections = []
        if self.yolo_model:
            results = self.yolo_model(pil_img, conf=0.25, verbose=False)
            
            for r in results:
                if r.boxes is not None:
                    for box in r.boxes:
                        class_id = int(box.cls[0])
                        class_name = self.yolo_model.names[class_id].lower()
                        confidence = float(box.conf[0])
                        x1, y1, x2, y2 = box.xyxy[0].tolist()
                        
                        detections.append({
                            "class": class_name,
                            "confidence": confidence,
                            "bbox": [int(x1), int(y1), int(x2), int(y2)],
                        })
        
        self.current_detections = detections
        
        # Draw bounding boxes on image
        draw = ImageDraw.Draw(pil_img)
        for det in detections:
            x1, y1, x2, y2 = det["bbox"]
            # Draw red box
            draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
            # Draw label background
            draw.rectangle([x1, y1-18, x1 + len(det["class"]) * 7 + 10, y1], fill="red")
            draw.text((x1 + 2, y1 - 15), det["class"], fill="white")
        
        return pil_img, detections
    
    def get_current_image_np(self):
        """Get current image as numpy array for Gradio"""
        img, _ = self.get_current_image_with_boxes()
        return np.array(img)
    
    def speak(self, text):
        """Convert text to speech"""
        try:
            tts = gTTS(text=text, lang="en", slow=False)
            audio_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3").name
            tts.save(audio_path)
            return audio_path
        except Exception as e:
            print(f"TTS error: {e}")
            return None
    
    def ask_new_question(self):
        """Generate a new question based on current detections"""
        if not self.current_detections:
            return None, "No objects detected. Try a different page."
        
        # Filter detections that match page targets
        valid_targets = self.page_targets.get(self.current_page, [])
        eligible = [d for d in self.current_detections if d["class"] in valid_targets]
        
        if not eligible:
            eligible = self.current_detections
        
        self.current_target = random.choice(eligible)
        question = f"Where is the {self.current_target['class']}? Click on it in the picture!"
        audio = self.speak(question)
        
        return question, audio
    
    def check_answer_from_click(self, evt: gr.EventData):
        """Check if clicked coordinates match any detected object"""
        if not self.current_target:
            return "⚠️ Click 'Ask Question' first!", None, 0, 0, 0
        
        # Get click coordinates from event
        if evt._data and "index" in evt._data:
            x, y = evt._data["index"][0], evt._data["index"][1]
        else:
            return "Please click on the image first!", None, 0, 0, 0
        
        # Find which object was clicked
        clicked_object = None
        for det in self.current_detections:
            x1, y1, x2, y2 = det["bbox"]
            if x1 <= x <= x2 and y1 <= y <= y2:
                clicked_object = det["class"]
                break
        
        if not clicked_object:
            return "Click on a detected object (with red box)!", None, 0, 0, 0
        
        # Check if correct
        is_correct = (clicked_object == self.current_target["class"])
        
        if is_correct:
            feedback = f"✅ Correct! That IS the {self.current_target['class']}! Great job!"
            score = 1
        else:
            feedback = f"❌ Not quite. You clicked on {clicked_object}, but I asked for the {self.current_target['class']}. Try again!"
            score = 0
        
        # Log trial
        trial = {
            "timestamp": datetime.now().isoformat(),
            "page": self.current_page,
            "asked_for": self.current_target["class"],
            "clicked_on": clicked_object,
            "click_coordinates": (x, y),
            "correct": is_correct,
            "score": score
        }
        self.session_data.append(trial)
        
        audio = self.speak(feedback)
        total_score = sum(t["score"] for t in self.session_data)
        trials = len(self.session_data)
        
        return feedback, audio, score, total_score, trials
    
    def next_page(self):
        """Move to next page"""
        if self.current_page < self.total_pages - 1:
            self.current_page += 1
        self.current_target = None
        return self.get_current_image_np(), f"Page {self.current_page + 1} / {self.total_pages}"
    
    def prev_page(self):
        """Move to previous page"""
        if self.current_page > 0:
            self.current_page -= 1
        self.current_target = None
        return self.get_current_image_np(), f"Page {self.current_page + 1} / {self.total_pages}"
    
    def get_session_report(self):
        """Generate session summary"""
        correct = sum(1 for t in self.session_data if t["correct"])
        total = len(self.session_data)
        accuracy = correct / total if total > 0 else 0
        
        return {
            "session_summary": {
                "total_trials": total,
                "correct_responses": correct,
                "accuracy": f"{accuracy:.1%}",
                "pages_covered": list(set(t["page"] for t in self.session_data))
            },
            "trials": self.session_data,
            "export_date": datetime.now().isoformat()
        }

# ============================================
# 2. INITIALIZE
# ============================================

tutor = ABATutor()

# ============================================
# 3. GRADIO INTERFACE
# ============================================

css = """
.main-container { max-width: 1400px; margin: auto; }
.page-card { background: white; border-radius: 20px; padding: 20px; box-shadow: 0 10px 40px rgba(0,0,0,0.1); }
.feedback-box { background: #f0f4ff; border-radius: 15px; padding: 15px; margin: 10px 0; }
.question-card { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 15px; padding: 20px; }
.click-hint { background: #e8f5e9; padding: 10px; border-radius: 10px; text-align: center; font-size: 14px; margin-top: 15px; }
"""

with gr.Blocks(title="ABA-AI-Tutor") as demo:
    gr.Markdown("""
    # 🖱️🤖 ABA-AI-Tutor
    
    ### Click-to-Select VB-MAPP Assessment
    *The AI asks "Where is the ___?" You **click on the object** in the picture to answer!*
    """)
    
    with gr.Row(elem_classes=["main-container"]):
        with gr.Column(scale=3):
            with gr.Group(elem_classes=["page-card"]):
                # Image with click event - THIS IS THE KEY!
                current_image = gr.Image(
                    value=tutor.get_current_image_np(),
                    label="Click on the object in this picture",
                    interactive=True,
                    height=500,
                    type="numpy"
                )
                with gr.Row():
                    prev_btn = gr.Button("◀ Previous Page", size="sm", variant="secondary")
                    page_info = gr.Textbox(value=f"Page {tutor.current_page + 1} / {tutor.total_pages}", interactive=False, container=False)
                    next_btn = gr.Button("Next Page ▶", size="sm", variant="secondary")
        
        with gr.Column(scale=2):
            gr.Markdown("### 🎯 Click Assessment")
            
            with gr.Group(elem_classes=["question-card"]):
                question_display = gr.Textbox(label="📢 AI Question", interactive=False, placeholder="Click 'Ask Question' to start...")
                question_audio = gr.Audio(label="🔊 Question Audio", type="filepath", interactive=False)
            
            with gr.Row():
                ask_btn = gr.Button("🎲 Ask New Question", variant="primary", size="lg")
            
            with gr.Group(elem_classes=["feedback-box"]):
                feedback_msg = gr.Textbox(label="💬 AI Feedback", interactive=False)
                feedback_audio = gr.Audio(label="🔊 Feedback Audio", type="filepath", interactive=False)
            
            with gr.Row():
                trial_score = gr.Number(label="⭐ Last Score", interactive=False, value=0)
                total_score_display = gr.Number(label="🏆 Total Score", interactive=False, value=0)
                trials_count = gr.Number(label="📋 Trials Completed", interactive=False, value=0)
            
            gr.Markdown("""
            <div class="click-hint">
            🖱️ <strong>How to play:</strong><br>
            1. Click "Ask New Question"<br>
            2. Listen to what the AI asks for<br>
            3. <strong>Click directly on that object</strong> in the picture<br>
            4. The AI automatically checks your answer when you click!<br>
            🎉 No extra button needed!
            </div>
            """)
    
    with gr.Tabs():
        with gr.TabItem("📊 Session Report"):
            refresh_btn = gr.Button("Refresh Report")
            report_json = gr.JSON(label="Full Assessment Data", value=tutor.get_session_report())
            accuracy_display = gr.Markdown("**Accuracy:** 0%")
    
    # ============================================
    # 4. EVENT HANDLERS
    # ============================================
    
    def ask_question():
        question, audio = tutor.ask_new_question()
        if question:
            return question, audio, "Waiting for your click...", None, 0, 0, 0
        return "No objects detected. Try a different page.", None, "⚠️ No clickable objects", None, 0, 0, 0
    
    ask_btn.click(
        ask_question,
        outputs=[question_display, question_audio, feedback_msg, feedback_audio, trial_score, total_score_display, trials_count]
    )
    
    # CRITICAL: This connects image clicks to the assessment
    def on_image_click(evt: gr.EventData):
        if not tutor.current_target:
            return "👂 Click 'Ask Question' first!", None, 0, 0, 0
        
        feedback, audio, score, total, trials = tutor.check_answer_from_click(evt)
        
        # Also refresh the image to show which object was clicked (optional)
        new_image = tutor.get_current_image_np()
        
        return new_image, feedback, audio, score, total, trials
    
    # Connect click event on the image
    current_image.select(
        fn=on_image_click,
        outputs=[current_image, feedback_msg, feedback_audio, trial_score, total_score_display, trials_count]
    )
    
    def prev_page_update():
        img, info = tutor.prev_page()
        return img, info, "", None, "", None, 0, 0, 0
    
    def next_page_update():
        img, info = tutor.next_page()
        return img, info, "", None, "", None, 0, 0, 0
    
    prev_btn.click(
        prev_page_update,
        outputs=[current_image, page_info, question_display, question_audio, feedback_msg, feedback_audio, trial_score, total_score_display, trials_count]
    )
    
    next_btn.click(
        next_page_update,
        outputs=[current_image, page_info, question_display, question_audio, feedback_msg, feedback_audio, trial_score, total_score_display, trials_count]
    )
    
    def refresh_report():
        report = tutor.get_session_report()
        accuracy = report["session_summary"]["accuracy"]
        return report, f"**Accuracy:** {accuracy}"
    
    refresh_btn.click(refresh_report, outputs=[report_json, accuracy_display])

# ============================================
# 5. LAUNCH
# ============================================

if __name__ == "__main__":
    demo.launch()