import gradio as gr
import torch
import clip
from PIL import Image
import numpy as np
import json
from typing import List, Tuple
import io
import base64

# Load CLIP model
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

# Load categories
def load_categories():
    """Load drawing categories from categories.txt file"""
    try:
        with open("categories.txt", "r") as f:
            categories = [line.strip() for line in f.readlines()]
        return categories
    except FileNotFoundError:
        # Fallback to a subset of common categories if file not found
        return [
            "airplane", "apple", "car", "cat", "dog", "flower", "house", "tree",
            "sun", "moon", "star", "bird", "fish", "butterfly", "bicycle", "boat"
        ]

# Load categories
CATEGORIES = load_categories()

def preprocess_drawing(image):
    """Preprocess the drawing image for CLIP"""
    if image is None:
        return None
    print(image)
    # Convert to PIL Image if it's not already
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    # print(image)
    # Convert to RGB if it's not already
    if image.mode != 'RGB':
        image = image.convert('RGB')
    # Save the image to a temporary file for debugging
    image.save("temp_drawing.png")
    # Preprocess for CLIP
    return preprocess(image).unsqueeze(0).to(device)

def classify_drawing(image, top_k: int = 5) -> List[Tuple[str, float]]:
    """
    Classify a quickdraw-style drawing using CLIP zero-shot classification
    
    Args:
        image: PIL Image or numpy array of the drawing
        top_k: Number of top predictions to return
    
    Returns:
        List of tuples (category, confidence_score)
    """
    if image is None:
        print("Here")
        return []
    
    try:
        # Preprocess the image
        if isinstance(image, dict):
            print("Image is a dictionary")
            # return []
            print(image)
            image_input = preprocess_drawing(image['composite'])
        else:
            image_input = preprocess_drawing(image)
        if image_input is None:
            # print("Here2")
            return []
        
        # Create text prompts for all categories
        text_prompts = [f"a drawing of a {category}" for category in CATEGORIES]
        text_inputs = clip.tokenize(text_prompts).to(device)
        
        # Get features
        with torch.no_grad():
            image_features = model.encode_image(image_input)
            text_features = model.encode_text(text_inputs)
            
            # Normalize features
            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
            
            # Calculate similarities
            similarities = (100.0 * image_features @ text_features.T).softmax(dim=-1)
            values, indices = similarities[0].topk(top_k)
        
        # Format results
        results = []
        for i in range(top_k):
            category = CATEGORIES[indices[i]]
            confidence = float(values[i])
            results.append((category, confidence))
        
        return results
    
    except Exception as e:
        print(f"Error in classification: {str(e)}")
        return []
    
def classify_drawing1(image, top_k: int = 5,word: str = "dog") -> List[Tuple[str, float]]:
    """
    Classify a quickdraw-style drawing using CLIP zero-shot classification
    
    Args:
        image: PIL Image or numpy array of the drawing
        top_k: Number of top predictions to return
    
    Returns:
        List of tuples (category, confidence_score)
    """
    if image is None:
        print("Here")
        return []
    
    try:
        # Preprocess the image
        if isinstance(image, dict):
            print("Image is a dictionary")
            # return []
            print(image)
            image_input = preprocess_drawing(image['composite'])
        else:
            image_input = preprocess_drawing(image)
        if image_input is None:
            # print("Here2")
            return []
        if word.lower() not in [cat.lower() for cat in CATEGORIES]:
            CATEGORIES.append(word)
            # print(f"Added new category: {word}")
        
        # Create text prompts for all categories
        text_prompts = [f"a drawing of a {category}" for category in CATEGORIES]
        text_inputs = clip.tokenize(text_prompts).to(device)
        
        # Get features
        with torch.no_grad():
            image_features = model.encode_image(image_input)
            text_features = model.encode_text(text_inputs)
            
            # Normalize features
            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
            
            # Calculate similarities
            similarities = (100.0 * image_features @ text_features.T).softmax(dim=-1)
            values, indices = similarities[0].topk(top_k)
        
        # Format results
        results = []
        for i in range(top_k):
            category = CATEGORIES[indices[i]]
            confidence = float(values[i])
            results.append((category, confidence))

        if word.lower() not in [cat.lower() for cat, _ in results]:
            # Find the original confidence score for the word
            word_idx = CATEGORIES.index(word)
            word_confidence = float(similarities[0][word_idx])
            
            # Remove the last result if we have top_k results
            if len(results) >= top_k:
                results.pop()
            # Add the word with its original confidence score
            results.append((word, word_confidence))
        
        return results
    
    except Exception as e:
        print(f"Error in classification: {str(e)}")
        return []

def gradio_classify(image):
    """Gradio interface function"""
    if image is None:
        return "Please draw something first!"
    
    results = classify_drawing(image, top_k=5)
    
    if not results:
        return "Sorry, couldn't classify the drawing. Please try again."
    
    # Format results for display
    output = "🎨 **Drawing Classification Results:**\n\n"
    for i, (category, confidence) in enumerate(results, 1):
        percentage = confidence * 100
        output += f"{i}. **{category.title()}** - {percentage:.1f}%\n"
    
    return output

# Create Gradio interface
def create_interface():
    """Create the Gradio interface with both drawing and API tabs"""
    
    with gr.Blocks(
        title="QuickDraw Classifier",
        theme=gr.themes.Soft(),
        css="""
        .main-header { text-align: center; margin-bottom: 2rem; }
        .description { text-align: center; margin-bottom: 1.5rem; font-size: 1.1rem; }
        """
    ) as demo:
        
        gr.HTML("""
        <div class="main-header">
            <h1>🎨 QuickDraw Classifier</h1>
            <p class="description">
                Draw anything and let AI guess what it is! This app uses CLIP for zero-shot classification 
                of your drawings across 345+ categories from Google's QuickDraw dataset.
            </p>
        </div>
        """)
        
        with gr.Tabs():
            # Main Drawing Interface Tab
            with gr.Tab("🖊️ Draw & Classify"):
                with gr.Row():
                    with gr.Column(scale=1):
                        # Drawing canvas
                        canvas = gr.Sketchpad(
                            label="Draw here! 🖊️",
                            type="pil",
                            image_mode="RGB",
                            canvas_size=(400, 400),
                            brush=gr.Brush(colors=["#000000"], color_mode="fixed", default_size=5),
                            show_download_button=False,
                            show_share_button=False,
                        )
                        
                        with gr.Row():
                            classify_btn = gr.Button("🔍 Classify Drawing", variant="primary", size="lg")
                            clear_btn = gr.ClearButton(canvas, value="🗑️ Clear Canvas", size="lg")
                    
                    with gr.Column(scale=1):
                        # Results
                        output = gr.Markdown(
                            label="Classification Results",
                            value="👈 Draw something and click 'Classify Drawing' to see results!"
                        )
                
                # Examples section
                gr.HTML("""
                <div style="margin-top: 2rem; text-align: center;">
                    <h3>💡 Drawing Tips:</h3>
                    <p>• Draw simple, clear shapes<br>
                    • Try common objects like animals, vehicles, or everyday items<br>
                    • The AI knows 345+ categories including: airplane, cat, house, tree, car, flower, etc.</p>
                </div>
                """)
                
                # Event handlers
                classify_btn.click(
                    fn=gradio_classify,
                    inputs=[canvas],
                    outputs=[output]
                )
                
                # Auto-classify on canvas change (optional, can be removed if too frequent)
                canvas.change(
                    fn=gradio_classify,
                    inputs=[canvas],
                    outputs=[output]
                )
            
            # API Interface Tab
            with gr.Tab("🔌 API Testing"):
                gr.HTML("""
                <div style="text-align: center; margin-bottom: 1rem;">
                    <h3>Base64 Image Classification API</h3>
                    <p>Test the API endpoint with base64 encoded images</p>
                </div>
                """)
                
                with gr.Row():
                    with gr.Column():
                        api_input = gr.Textbox(
                            label="Base64 Image Data", 
                            placeholder="Paste base64 encoded image here...",
                            lines=3
                        )
                        api_top_k = gr.Slider(
                            minimum=1, 
                            maximum=10, 
                            value=5, 
                            step=1, 
                            label="Top K Predictions"
                        )
                        api_word = gr.Textbox(
                            label="Word to Add",
                            placeholder="Enter a word to add to the categories",
                            lines=1
                        )
                        api_btn = gr.Button("🚀 Classify", variant="primary")
                    
                    with gr.Column():
                        api_output = gr.JSON(label="API Response")
                
                api_btn.click(
                    fn=classify_image_api,
                    inputs=[api_input, api_top_k, api_word],
                    outputs=[api_output]
                )
    
    return demo

# API endpoint function for external use
def classify_image_api(image_data: str, top_k: int = 5, word: str = "dog"):
    """
    API endpoint for classifying base64 encoded images
    
    Args:
        image_data: Base64 encoded image string
        top_k: Number of top predictions to return
    
    Returns:
        JSON response with classification results
    """
    try:
        # Decode base64 image
        image_bytes = base64.b64decode(image_data)
        image = Image.open(io.BytesIO(image_bytes))
        
        # Classify
        results = classify_drawing1(image, top_k, word)
        # Normalize confidence scores to sum to 100%
        total_conf = sum(conf for _, conf in results)
        if total_conf > 0:
            results = [(cat, (conf/total_conf) * 100) for cat, conf in results]
        
        return {
            "success": True,
            "predictions": [
                {"category": cat, "confidence": conf} 
                for cat, conf in results
            ]
        }
    
    except Exception as e:
        return {
            "success": False,
            "error": str(e)
        }

if __name__ == "__main__":
    # Create and launch the main interface
    demo = create_interface()
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )