import gradio as gr
import os
import torch
from model import create_model
from timeit import default_timer as timer
from typing import Tuple, Dict

# Setup class names
class_names = ['Fake', 'Real']

### 2. Model and transforms preparation ###
# Create model
model, transforms = create_model(num_classes=2)

# Load saved weights
model.load_state_dict(
    torch.load(
        f="RealityCheck.pth",
        map_location=torch.device("cpu"),
    )
)

### 3. Predict function ###
def predict(img) -> Tuple[Dict, float]:
    """Transforms and performs a prediction on img and returns prediction and time taken."""
    # Start the timer
    start_time = timer()
    
    # Transform the target image and add a batch dimension
    img = transforms(img).unsqueeze(0)
    
    # Put model into evaluation mode and turn on inference mode
    model.eval()
    with torch.inference_mode():
        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
        pred_probs = torch.softmax(model(img), dim=1)
    
    # Create a prediction label and prediction probability dictionary for each prediction class
    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
    
    # Calculate the prediction time
    pred_time = round(timer() - start_time, 5)
    
    return pred_labels_and_probs, pred_time

### 4. Enhanced Gradio app ###
# Custom CSS for ultra-modern, attractive styling
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

* {
    font-family: 'Inter', sans-serif;
}

body {
    background: linear-gradient(-45deg, #ee7752, #e73c7e, #23a6d5, #23d5ab);
    background-size: 400% 400%;
    animation: gradientBG 15s ease infinite;
    margin: 0;
    padding: 0;
}

@keyframes gradientBG {
    0% { background-position: 0% 50%; }
    50% { background-position: 100% 50%; }
    100% { background-position: 0% 50%; }
}

.gradio-container {
    background: transparent !important;
    min-height: 100vh;
    position: relative;
    overflow-x: hidden;
}

.gradio-container::before {
    content: '';
    position: fixed;
    top: 0;
    left: 0;
    right: 0;
    bottom: 0;
    background: 
        radial-gradient(circle at 20% 50%, rgba(120, 119, 198, 0.3) 0%, transparent 50%),
        radial-gradient(circle at 80% 20%, rgba(255, 119, 198, 0.3) 0%, transparent 50%),
        radial-gradient(circle at 40% 80%, rgba(120, 219, 255, 0.3) 0%, transparent 50%);
    pointer-events: none;
    z-index: -1;
}

#main-container {
    max-width: 1200px !important;
    margin: 0 auto;
    padding: 20px;
    position: relative;
    z-index: 1;
}

#title {
    text-align: center;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-size: 4em;
    font-weight: 700;
    margin-bottom: 10px;
    text-shadow: 0 0 30px rgba(102, 126, 234, 0.5);
    animation: titleGlow 3s ease-in-out infinite alternate;
}

@keyframes titleGlow {
    from { filter: drop-shadow(0 0 20px rgba(102, 126, 234, 0.3)); }
    to { filter: drop-shadow(0 0 40px rgba(102, 126, 234, 0.8)); }
}

#subtitle {
    text-align: center;
    color: rgba(255,255,255,0.95);
    font-size: 1.4em;
    font-weight: 300;
    margin-bottom: 40px;
    text-shadow: 0 2px 10px rgba(0,0,0,0.4);
    animation: fadeInUp 1s ease-out;
}

@keyframes fadeInUp {
    from { opacity: 0; transform: translateY(30px); }
    to { opacity: 1; transform: translateY(0); }
}

.stats-container {
    display: flex;
    justify-content: space-around;
    margin: 40px 0;
    gap: 20px;
    animation: slideInUp 1.2s ease-out;
}

@keyframes slideInUp {
    from { opacity: 0; transform: translateY(50px); }
    to { opacity: 1; transform: translateY(0); }
}

.stat-item {
    background: rgba(255,255,255,0.15);
    backdrop-filter: blur(20px);
    border: 1px solid rgba(255,255,255,0.2);
    border-radius: 20px;
    padding: 25px 20px;
    color: white;
    text-align: center;
    flex: 1;
    transition: all 0.4s cubic-bezier(0.4, 0, 0.2, 1);
    position: relative;
    overflow: hidden;
}

.stat-item::before {
    content: '';
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255,255,255,0.2), transparent);
    transition: left 0.8s;
}

.stat-item:hover::before {
    left: 100%;
}

.stat-item:hover {
    transform: translateY(-10px) scale(1.05);
    box-shadow: 0 20px 40px rgba(0,0,0,0.3);
    background: rgba(255,255,255,0.25);
}

.stat-number {
    font-size: 2.5em;
    font-weight: 700;
    background: linear-gradient(135deg, #fff, #e0e0e0);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin-bottom: 5px;
}

.guide-box {
    background: rgba(255,255,255,0.12);
    backdrop-filter: blur(20px);
    border: 1px solid rgba(255,255,255,0.2);
    border-radius: 25px;
    padding: 35px;
    margin: 30px 0;
    box-shadow: 0 15px 35px rgba(0,0,0,0.1);
    animation: floatIn 1.5s ease-out;
    position: relative;
    overflow: hidden;
}

@keyframes floatIn {
    from { opacity: 0; transform: translateY(30px) scale(0.95); }
    to { opacity: 1; transform: translateY(0) scale(1); }
}

.guide-box::before {
    content: '';
    position: absolute;
    top: -50%;
    left: -50%;
    width: 200%;
    height: 200%;
    background: radial-gradient(circle, rgba(255,255,255,0.1) 0%, transparent 70%);
    animation: rotate 20s linear infinite;
    pointer-events: none;
}

@keyframes rotate {
    from { transform: rotate(0deg); }
    to { transform: rotate(360deg); }
}

.guide-box h3 {
    color: white;
    margin-bottom: 20px;
    font-size: 1.5em;
    font-weight: 600;
    text-shadow: 0 2px 10px rgba(0,0,0,0.3);
}

.guide-box ul {
    color: rgba(255,255,255,0.9);
    line-height: 1.8;
    font-size: 1.1em;
}

.guide-box li {
    margin-bottom: 12px;
    transition: transform 0.2s ease;
}

.guide-box li:hover {
    transform: translateX(10px);
}

.main-interface {
    background: rgba(255,255,255,0.08);
    backdrop-filter: blur(25px);
    border: 1px solid rgba(255,255,255,0.15);
    border-radius: 30px;
    padding: 40px;
    margin: 40px 0;
    box-shadow: 0 25px 50px rgba(0,0,0,0.2);
    animation: materialIn 1.8s ease-out;
}

@keyframes materialIn {
    from { opacity: 0; transform: scale(0.9) translateY(50px); }
    to { opacity: 1; transform: scale(1) translateY(0); }
}

.gr-button {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    border-radius: 15px !important;
    padding: 15px 40px !important;
    color: white !important;
    font-weight: 600 !important;
    font-size: 1.1em !important;
    cursor: pointer !important;
    transition: all 0.4s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow: 0 10px 25px rgba(102, 126, 234, 0.3) !important;
    position: relative !important;
    overflow: hidden !important;
}

.gr-button::before {
    content: '';
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255,255,255,0.3), transparent);
    transition: left 0.6s;
}

.gr-button:hover::before {
    left: 100%;
}

.gr-button:hover {
    transform: translateY(-3px) !important;
    box-shadow: 0 15px 35px rgba(102, 126, 234, 0.4) !important;
    background: linear-gradient(135deg, #5a6fd8 0%, #6a4c93 100%) !important;
}

.gr-button:active {
    transform: translateY(0px) !important;
}

.gr-form {
    background: rgba(255,255,255,0.05) !important;
    backdrop-filter: blur(15px) !important;
    border: 1px solid rgba(255,255,255,0.1) !important;
    border-radius: 20px !important;
    padding: 25px !important;
}

.gr-input, .gr-textbox, .gr-dropdown {
    background: rgba(255,255,255,0.1) !important;
    border: 1px solid rgba(255,255,255,0.2) !important;
    border-radius: 12px !important;
    color: white !important;
    backdrop-filter: blur(10px) !important;
}

.gr-input:focus, .gr-textbox:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 20px rgba(102, 126, 234, 0.3) !important;
}

.footer {
    text-align: center;
    color: rgba(255,255,255,0.8);
    margin-top: 50px;
    padding: 30px;
    font-size: 1.1em;
    animation: fadeIn 2s ease-out;
}

@keyframes fadeIn {
    from { opacity: 0; }
    to { opacity: 1; }
}

.footer::before {
    content: '✨';
    display: inline-block;
    margin-right: 10px;
    animation: sparkle 2s ease-in-out infinite;
}

@keyframes sparkle {
    0%, 100% { transform: scale(1) rotate(0deg); }
    50% { transform: scale(1.2) rotate(180deg); }
}

/* Responsive Design */
@media (max-width: 768px) {
    #title { font-size: 2.5em; }
    .stats-container { flex-direction: column; gap: 15px; }
    .main-interface { padding: 25px; margin: 20px 0; }
    .guide-box { padding: 25px; }
}

/* Enhanced Gradio Component Styling */
.gr-box, .gr-form, .gr-panel {
    background: rgba(30, 30, 30, 0.95) !important;
    backdrop-filter: blur(20px) !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    border-radius: 20px !important;
    padding: 25px !important;
    box-shadow: 0 20px 40px rgba(0, 0, 0, 0.3) !important;
    transition: all 0.4s cubic-bezier(0.4, 0, 0.2, 1) !important;
    position: relative !important;
    overflow: hidden !important;
}

.gr-box::before, .gr-form::before, .gr-panel::before {
    content: '';
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.1), transparent);
    transition: left 0.8s;
    pointer-events: none;
}

.gr-box:hover::before, .gr-form:hover::before, .gr-panel:hover::before {
    left: 100%;
}

.gr-box:hover, .gr-form:hover, .gr-panel:hover {
    transform: translateY(-5px) !important;
    box-shadow: 0 25px 50px rgba(0, 0, 0, 0.4) !important;
    border-color: rgba(255, 255, 255, 0.2) !important;
}

/* Image Upload Container */
div[data-testid="image"] {
    background: rgba(20, 20, 20, 0.9) !important;
    border: 2px dashed rgba(255, 255, 255, 0.3) !important;
    border-radius: 20px !important;
    padding: 30px !important;
    transition: all 0.4s ease !important;
    position: relative !important;
    overflow: hidden !important;
}

div[data-testid="image"]:hover {
    border-color: #667eea !important;
    background: rgba(25, 25, 25, 0.95) !important;
    transform: scale(1.02) !important;
    box-shadow: 0 15px 30px rgba(102, 126, 234, 0.2) !important;
}

div[data-testid="image"]::before {
    content: '📸 Drop your image here or click to browse';
    position: absolute;
    top: 50%;
    left: 50%;
    transform: translate(-50%, -50%);
    color: rgba(255, 255, 255, 0.6);
    font-size: 1.2em;
    font-weight: 500;
    text-align: center;
    pointer-events: none;
    opacity: 0;
    transition: opacity 0.3s ease;
}

div[data-testid="image"]:empty::before {
    opacity: 1;
}

/* Enhanced Button Styling */
.gr-button {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    border-radius: 20px !important;
    padding: 18px 50px !important;
    color: white !important;
    font-weight: 600 !important;
    font-size: 1.2em !important;
    cursor: pointer !important;
    transition: all 0.4s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow: 0 15px 35px rgba(102, 126, 234, 0.4) !important;
    position: relative !important;
    overflow: hidden !important;
    text-transform: uppercase !important;
    letter-spacing: 1px !important;
    margin: 20px 0 !important;
}

.gr-button::before {
    content: '';
    position: absolute;
    top: 50%;
    left: 50%;
    width: 0;
    height: 0;
    background: rgba(255, 255, 255, 0.2);
    border-radius: 50%;
    transform: translate(-50%, -50%);
    transition: width 0.6s, height 0.6s;
}

.gr-button:hover::before {
    width: 300px;
    height: 300px;
}

.gr-button:hover {
    transform: translateY(-8px) scale(1.05) !important;
    box-shadow: 0 25px 50px rgba(102, 126, 234, 0.6) !important;
    background: linear-gradient(135deg, #5a6fd8 0%, #6a4c93 100%) !important;
}

.gr-button:active {
    transform: translateY(-4px) scale(1.02) !important;
    transition: all 0.1s ease !important;
}

/* Pulse Animation for Button */
@keyframes buttonPulse {
    0% { box-shadow: 0 15px 35px rgba(102, 126, 234, 0.4); }
    50% { box-shadow: 0 15px 35px rgba(102, 126, 234, 0.7); }
    100% { box-shadow: 0 15px 35px rgba(102, 126, 234, 0.4); }
}

.gr-button:not(:hover) {
    animation: buttonPulse 3s ease-in-out infinite;
}

/* Label and Output Styling */
.gr-label {
    background: rgba(20, 20, 20, 0.9) !important;
    border: 1px solid rgba(255, 255, 255, 0.15) !important;
    border-radius: 15px !important;
    padding: 20px !important;
    margin: 10px 0 !important;
    color: white !important;
    backdrop-filter: blur(15px) !important;
    transition: all 0.3s ease !important;
}

.gr-label:hover {
    background: rgba(25, 25, 25, 0.95) !important;
    transform: translateY(-2px) !important;
    box-shadow: 0 10px 20px rgba(0, 0, 0, 0.3) !important;
}

/* Progress Bars */
.gr-label .bar {
    background: linear-gradient(90deg, #667eea, #764ba2) !important;
    border-radius: 10px !important;
    height: 8px !important;
    position: relative !important;
    overflow: hidden !important;
}

.gr-label .bar::before {
    content: '';
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.4), transparent);
    animation: shimmer 2s infinite;
}

@keyframes shimmer {
    0% { left: -100%; }
    100% { left: 100%; }
}

/* Number Input Styling */
.gr-number input {
    background: rgba(20, 20, 20, 0.9) !important;
    border: 1px solid rgba(255, 255, 255, 0.2) !important;
    border-radius: 12px !important;
    color: white !important;
    padding: 15px !important;
    font-size: 1.1em !important;
    font-weight: 500 !important;
    text-align: center !important;
    backdrop-filter: blur(10px) !important;
    transition: all 0.3s ease !important;
}

.gr-number input:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 20px rgba(102, 126, 234, 0.3) !important;
    transform: scale(1.02) !important;
}

/* Examples Gallery Styling */
.gr-gallery {
    background: rgba(20, 20, 20, 0.8) !important;
    border-radius: 20px !important;
    padding: 20px !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    backdrop-filter: blur(15px) !important;
}

.gr-gallery img {
    border-radius: 15px !important;
    transition: all 0.3s ease !important;
    border: 2px solid rgba(255, 255, 255, 0.1) !important;
}

.gr-gallery img:hover {
    transform: scale(1.1) rotate(2deg) !important;
    border-color: #667eea !important;
    box-shadow: 0 10px 25px rgba(102, 126, 234, 0.4) !important;
}

/* Loading Animation */
.loading-overlay {
    position: absolute;
    top: 0;
    left: 0;
    right: 0;
    bottom: 0;
    background: rgba(0, 0, 0, 0.8);
    display: flex;
    align-items: center;
    justify-content: center;
    border-radius: 20px;
    z-index: 1000;
}

.loading-spinner {
    width: 50px;
    height: 50px;
    border: 4px solid rgba(255, 255, 255, 0.2);
    border-left: 4px solid #667eea;
    border-radius: 50%;
    animation: spin 1s linear infinite;
}

/* Animated Background Patterns */
.pattern-overlay {
    position: absolute;
    top: 0;
    left: 0;
    right: 0;
    bottom: 0;
    background-image: 
        radial-gradient(circle at 25% 25%, rgba(102, 126, 234, 0.1) 0%, transparent 50%),
        radial-gradient(circle at 75% 75%, rgba(118, 75, 162, 0.1) 0%, transparent 50%);
    pointer-events: none;
    animation: patternMove 20s ease-in-out infinite;
}

@keyframes patternMove {
    0%, 100% { transform: translate(0, 0) rotate(0deg); }
    50% { transform: translate(10px, -10px) rotate(180deg); }
}

/* Responsive Enhancements */
@media (max-width: 768px) {
    .gr-button {
        padding: 14px 30px !important;
        font-size: 1em !important;
    }
    
    .gr-box, .gr-form, .gr-panel {
        padding: 20px !important;
        border-radius: 15px !important;
    }
    
    div[data-testid="image"] {
        padding: 20px !important;
    }
}
"""

# Create compact title and description
title = "🔍 AI-Image Detector"
description = """
<div id="main-container">
    <h1 id="title">🔍 AI-Image Detector</h1>
    <p id="subtitle">AI-Powered Real vs Fake Faces Detection</p>
    
    <div class="stats-container">
        <div class="stat-item">
            <div class="stat-number">95%</div>
            <div>Accuracy Rate</div>
        </div>
        <div class="stat-item">
            <div class="stat-number">&lt;1s</div>
            <div>Lightning Fast</div>
        </div>
        <div class="stat-item">
            <div class="stat-number">🤖</div>
            <div>AI Powered</div>
        </div>
        <div class="stat-item">
            <div class="stat-number">∞</div>
            <div>Unlimited Use</div>
        </div>
    </div>
    
    <div class="guide-box">
        <h3>🚀 How It Works</h3>
        <ul>
            <li><strong>🎯 Advanced AI Detection:</strong> Our neural network analyzes facial features, textures, and patterns to detect AI-generated content</li>
            <li><strong>📸 Simply Upload:</strong> Drag & drop or click to upload any face image - supports JPG, PNG, WebP formats</li>
            <li><strong>⚡ Real-time Analysis:</strong> Get instant results with confidence scores and processing time</li>
            <li><strong>🔒 Privacy First:</strong> Your images are processed securely and never stored on our servers</li>
        </ul>
    </div>
</div>
"""

# Create examples list from "examples/" directory
try:
    example_list = [["examples/" + example] for example in os.listdir("examples")] if os.path.exists("examples") else []
except:
    example_list = []

# Create the enhanced Gradio interface
with gr.Blocks(css=custom_css, title="Reality Check - AI Face Detector") as demo:
    gr.HTML(description)
    
    with gr.Row(elem_classes="main-interface"):
        with gr.Column(scale=1):
            input_image = gr.Image(
                type="pil",
                label="📤 Upload Face Image",
                height=400
            )
            
            if example_list:
                gr.Examples(
                    examples=example_list,
                    inputs=input_image,
                    label="🖼️ Try These Examples"
                )
        
        with gr.Column(scale=1):
            with gr.Group():
                output_label = gr.Label(
                    num_top_classes=2,
                    label="🎯 Classification Results",
                    show_label=True
                )
                output_time = gr.Number(
                    label="⚡ Processing Time (seconds)",
                    show_label=True
                )
                
                # Additional info section
                gr.HTML("""
                <div style="margin-top: 20px; padding: 20px; background: rgba(255,255,255,0.1); border-radius: 15px; backdrop-filter: blur(10px);">
                    <h4 style="color: white; margin-bottom: 15px;">🔍 Detection Insights</h4>
                    <div style="color: rgba(255,255,255,0.9); font-size: 0.9em; line-height: 1.6;">
                        • <strong>Real faces</strong> show natural imperfections and organic patterns<br>
                        • <strong>AI faces</strong> often have subtle artifacts in eyes, teeth, or background<br>
                        • Higher confidence scores indicate more certain predictions
                    </div>
                </div>
                """)
    
    # Submit button
    submit_btn = gr.Button(
    value="✨ 🔍 Analyze Image ✨",
    variant="primary",
    scale=3,
    min_width=280,
    interactive=True
    )


    # Connect the function
    submit_btn.click(
        fn=predict,
        inputs=input_image,
        outputs=[output_label, output_time]
    )
    
    # Auto-submit when image is uploaded
    input_image.change(
        fn=predict,
        inputs=input_image,
        outputs=[output_label, output_time],
        show_progress=True
    )
    
    # Footer
    gr.HTML("""
    <div class="footer">
        <p>🚀 Powered by Advanced AI Technology | 🔒 Privacy-First Design | ⚡ Lightning Fast Results</p>
        <div style="margin-top: 15px; font-size: 0.9em; opacity: 0.8;">
            Built with ❤️ for the AI community | By Nizamuddin Mandekar
        </div>
    </div>
    """)

# Launch the demo
if __name__ == "__main__":
    demo.launch(
        share=True,
        inbrowser=True,
        show_error=True,
        favicon_path=None,  # You can add your own favicon
        app_kwargs={"docs_url": None, "redoc_url": None}
    )