"""
AI Image Detector - Hugging Face Spaces Deployment
Detects whether an image is AI-generated or real using a MobileNetV2-based model.
"""

import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input

# Load the trained model
def load_model():
    """Load the trained transfer learning model."""
    model = tf.keras.models.load_model(
        "transfer_model.keras",
        custom_objects={'preprocess_input': preprocess_input},
    )
    return model

# Initialize model (loaded once at startup)
model = load_model()


def preprocess_image(image):
    """Preprocess the input image for model inference."""
    # Convert to PIL Image if needed
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)

    # Convert to RGB if necessary
    if image.mode != 'RGB':
        image = image.convert('RGB')

    # Resize to model's expected input size
    image = image.resize((224, 224))

    # Convert to numpy array
    img_array = np.array(image).astype('float32')

    # Add batch dimension
    img_array = np.expand_dims(img_array, axis=0)

    # Note: preprocessing is handled by the Lambda layer in the model
    # So we don't need to apply preprocess_input here

    return img_array

def predict(image):
    """
    Predict whether an image is AI-generated or real.

    Args:
        image: Input image (PIL Image or numpy array)

    Returns:
        Dictionary with class labels and their probabilities
    """
    if model is None:
        return {"Error": 1.0}

    if image is None:
        return {"Error": 1.0}

    try:
        # Preprocess the image
        processed_image = preprocess_image(image)

        # Make prediction (removed training=False parameter - not supported in Keras 3)
        prediction = model.predict(processed_image, verbose=0)[0][0]

        # Convert to probabilities for each class
        ai_prob = float(prediction)
        real_prob = 1.0 - ai_prob

        return {
            "Real Image": real_prob,
            "AI-Generated": ai_prob
        }

    except Exception as e:
        print(f"Prediction error: {e}")
        return {"Error": 1.0}

def classify_image(image):
    """
    Main classification function with detailed output.

    Args:
        image: Input image

    Returns:
        Tuple of (label_output, confidence_text)
    """
    if model is None:
        return None, "❌ **Model failed to load. Please check the model file.**"

    if image is None:
        return None, "Please upload an image to analyze."

    results = predict(image)

    if "Error" in results:
        return None, "Error processing image. Please try again."

    # Determine the classification
    ai_prob = results["AI-Generated"]
    real_prob = results["Real Image"]

    if ai_prob > 0.5:
        classification = "AI-Generated"
        confidence = ai_prob
    else:
        classification = "Real Image"
        confidence = real_prob

    # Create detailed analysis text
    analysis = f"""
### Analysis Results
**Classification:** {classification}
**Confidence:** {confidence:.1%}
---
**Detailed Probabilities:**
- Real Image: {real_prob:.2%}
- AI-Generated: {ai_prob:.2%}
---
*Note: This model was trained on CIFAKE and Tiny GenImage datasets.
Results may vary for images outside these domains.*
"""

    return results, analysis

# Create the Gradio interface
with gr.Blocks(
    title="AI Image Detector"
) as demo:
    gr.Markdown(
        """
        # 🔍 AI Image Detector

        Upload an image to detect whether it's **AI-generated** or a **real photograph**.

        This model uses transfer learning with MobileNetV2 and was trained on the
        CIFAKE and Tiny GenImage datasets, achieving ~95% accuracy on validation data.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Upload Image",
                type="pil",
                height=400
            )

            submit_btn = gr.Button(
                "🔎 Analyze Image",
                variant="primary",
                size="lg"
            )

            gr.Markdown(
                """
                ### Tips for best results:
                - Use clear, high-quality images
                - The model works best on photographs and AI-generated images
                - Works with JPG, PNG, and other common formats
                """
            )

        with gr.Column(scale=1):
            output_label = gr.Label(
                label="Classification Results",
                num_top_classes=2
            )

            output_text = gr.Markdown(
                label="Detailed Analysis"
            )

    # Connect the interface
    submit_btn.click(
        fn=classify_image,
        inputs=input_image,
        outputs=[output_label, output_text]
    )

    input_image.change(
        fn=classify_image,
        inputs=input_image,
        outputs=[output_label, output_text]
    )

    gr.Markdown(
        """
        ---
        ### About this model

        - **Architecture:** MobileNetV2 with custom classification head
        - **Training Data:** CIFAKE dataset + Tiny GenImage dataset (~128,000 images)
        - **Input Size:** 224×224 pixels
        - **Classes:** Real (0) vs AI-Generated (1)

        **Disclaimer:** This tool is for educational and research purposes.
        AI detection is an evolving field, and no detector is 100% accurate.

        """
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()