# Deepfake Detection Gradio App - v1.1
import gradio as gr
import os
import sys
import json
import argparse
from types import SimpleNamespace
from PIL import Image
import matplotlib.pyplot as plt
import io
import numpy as np

# Try to import detector - if this fails, we'll show an error in the UI
try:
    from support.detect import run_detect
    DETECTOR_AVAILABLE = True
    IMPORT_ERROR = None
except Exception as e:
    DETECTOR_AVAILABLE = False
    IMPORT_ERROR = str(e)
    print(f"Warning: Could not import detector: {e}")
    # Create a dummy function
    def run_detect(args):
        raise ImportError(f"Detector not available: {IMPORT_ERROR}")

# Download weights on first run (for HF Spaces)
if os.environ.get("SPACE_ID"):
    try:
        from download_weights import download_all_weights
        download_all_weights()
    except Exception as e:
        print(f"Warning: Could not download weights: {e}")

# Available detectors based on launcher.py
DETECTORS = ['ALL', 'R50_TF', 'R50_nodown', 'CLIP-D', 'P2G', 'NPR', 'PGC', 'ForensicConcept', 'MDMF']


def process_image(image_path):
    """
    Check if image is larger than 1024x1024 and central crop it if necessary.
    Returns the path to the processed image (or original if no change).
    """
    try:
        with Image.open(image_path) as img:
            width, height = img.size
            
            # Check if both dimensions are larger than 1024
            if width > 1024 and height > 1024:
                print(f"Image size {width}x{height} exceeds 1024x1024. Performing central crop.")
                
                # Calculate crop box
                left = (width - 1024) / 2
                top = (height - 1024) / 2
                right = (width + 1024) / 2
                bottom = (height + 1024) / 2
                
                # Crop
                img_cropped = img.crop((left, top, right, bottom))
                
                # Save to new path
                directory, filename = os.path.split(image_path)
                name, ext = os.path.splitext(filename)
                new_filename = f"{name}_cropped{ext}"
                new_path = os.path.join(directory, new_filename)
                
                img_cropped.save(new_path)
                return new_path
                
            return image_path
    except Exception as e:
        print(f"Error processing image: {e}")
        return image_path

def run_single_detection(image_path, detector_name):
    output_path = f"temp_result_{detector_name}.json"
    # Mock args object
    args = SimpleNamespace(
        image=image_path,
        detector=detector_name,
        config_dir='configs',
        output=output_path,
        weights='pretrained', # Use default/pretrained
        device='cpu', # Force CPU
        batch_size=1,
        dry_run=False,
        verbose=False
    )
    
    try:
        run_detect(args)
        
        if os.path.exists(output_path):
            with open(output_path, 'r') as f:
                result = json.load(f)
            os.remove(output_path)
            return result
        return None
    except Exception as e:
        if os.path.exists(output_path):
            try:
                os.remove(output_path)
            except:
                pass
        print(f"Error running {detector_name}: {e}")
        return None

def predict(image_path, detector_name):
    # Check if detector is available
    if not DETECTOR_AVAILABLE:
        return json.dumps({
            "error": "Detector module not available",
            "details": IMPORT_ERROR,
            "message": "The detection system could not be initialized. Please check the logs."
        }, indent=2), None
    
    if not image_path:
        return json.dumps({"error": "Please upload an image."}, indent=2), None
    
    # Process image (central crop if too large)
    processed_path = image_path
    try:
        processed_path = process_image(image_path)
    except Exception as e:
        print(f"Warning: Image processing failed: {e}")
        # Continue with original image if processing fails
    
    try:
        if detector_name == 'ALL':
            results = []
            # Filter out 'ALL' from detectors list
            real_detectors = [d for d in DETECTORS if d != 'ALL']
            
            for det in real_detectors:
                res = run_single_detection(processed_path, det)
                if res:
                    results.append((det, res))
            
            if not results:
                return "Error: No results obtained from detectors.", None
            
            confidences = []
            labels = []
            colors = []
            lines = []

            for det, res in results:
                pred = res.get('prediction', 'Unknown')
                # fake_score = probability of fake for every detector (support/detect.py fills it in)
                raw_conf = res.get('fake_score', res.get('confidence', 0.0))

                # Display confidence (confidence of the prediction)
                if pred == 'fake':
                    score = raw_conf
                    color = 'red'
                else:
                    score = 1 - raw_conf
                    color = 'green'

                labels.append(det)
                confidences.append(score)
                colors.append(color)
                lines.append(f"{det}: {pred} (confidence {score:.4f})")

            explanation = "Results of each detector:\n" + "\n".join(lines)

            # Plotting
            fig, ax = plt.subplots(figsize=(10, 5))
            bars = ax.bar(labels, confidences, color=colors)
            ax.set_ylim(0, 1.05)
            ax.set_ylabel('Confidence')
            ax.set_title('Detector Confidence Scores')
            
            # Add custom legend for colors
            from matplotlib.patches import Patch
            legend_elements = [
                Patch(facecolor='green', label='Real'),
                Patch(facecolor='red', label='Fake')
            ]
            ax.legend(handles=legend_elements)
            
            # Add value labels
            for bar in bars:
                height = bar.get_height()
                ax.text(bar.get_x() + bar.get_width()/2., height,
                        f'{height:.2f}',
                        ha='center', va='bottom')
            
            plt.tight_layout()
            return explanation, fig
            
        else:
            # Single Detector
            res = run_single_detection(processed_path, detector_name)
            
            if res:
                prediction = res.get('prediction', 'Unknown')
                confidence = res.get('fake_score', res.get('confidence', 0.0))
                elapsed_time = res.get('elapsed_time', 0.0)
                
                if prediction == 'fake':  
                    output = {
                        "Prediction": prediction,
                        "Confidence": f"{confidence:.4f}",
                        "Elapsed Time": f"{elapsed_time:.3f}s"
                    }
                else:
                    output = {
                        "Prediction": prediction,
                        "Confidence": f"{1-confidence:.4f}",
                        "Elapsed Time": f"{elapsed_time:.3f}s"
                    }
                return json.dumps(output, indent=2), None
            else:
                return json.dumps({"error": "Detection failed"}), None

    except Exception as e:
        return json.dumps({"error": str(e)}), None
        
    finally:
        # Cleanup cropped image if it's different from original
        if processed_path != image_path and os.path.exists(processed_path):
            try:
                os.remove(processed_path)
            except Exception as e:
                print(f"Warning: Could not remove temporary file {processed_path}: {e}")

# Create Gradio Interface
# Use theme only if gradio version supports it
demo = gr.Blocks(title="Deepfake Detection Space", theme=gr.themes.Soft())

with demo:
    gr.Markdown("""
    # 🔍 Deepfake Detection Space
    This space provides an **educational, open-source demonstrator** collecting a series of baseline state-of-the-art methods for deepfake detection, allowing for free and unlimited use!

    If you want to train the detectors in this Space yourself, check out our [GitHub repository](https://github.com/truebees-ai/Image-Deepfake-Detectors-Public-Library).    
    
    ### ⚠️ Important
    * This tool showcases open-source detectors. It **does not** reflect the performance, robustness, or accuracy of Truebees' commercial detection methods. Want to give it a try? check **[Truebees Website](https://www.truebees.eu)**
    * These models were trained on image not shared on social networks from the **[DeepShield dataset](https://zenodo.org/records/15648378)**. You may experience a drop in accuracy when testing highly compressed images downloaded from social media platforms or messaging apps.
    * You can expect performance comparable to the results shown in [Dell'Anna et al. (2025)](https://arxiv.org/pdf/2504.20658).
    
    ⚠️ Due to file size limitations, model weights need to be downloaded automatically on first use. This may take a few moments. <br>
    ⚠️ To provide a free service, all models run on CPU. The detection process may take a few seconds, depending on the image size and the selected detector.
    """)
    with gr.Accordion("📊 Understanding the Results", open=False):
            gr.Markdown("""
            ### Individual Detectors
            * **Prediction**: Tells if an image is **Real** or **Fake**.
            * **Confidence**: The confidence with which the model determines if the image is real or fake.
            * **Elapsed Time**: The time the model needed to make the prediction (excluding preprocessing or model building).

            ### The "ALL" Option
            * Runs all available detectors (R50_TF, R50_nodown, CLIP-D, P2G, NPR, PGC, ForensicConcept, MDMF) sequentially on the input image.
            * Lists the prediction and confidence of each detector and generates a **Confidence Plot** visualizing each model's score. No aggregated verdict is produced.
            * In the plot, **Green** bars indicate a **Real** prediction, while **Red** bars indicate a **Fake** prediction.
            """)
    
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Input Image", height=400)
            detector_input = gr.Dropdown(
                choices=DETECTORS, 
                value=DETECTORS[0], 
                label="Select Detector",
                info="Choose which deepfake detection model to use"
            )
            submit_btn = gr.Button("🔍 Detect", variant="primary")
        
        with gr.Column():
            output_display = gr.Textbox(
                label="Detection Results", 
                lines=15, 
                max_lines=20,
                show_copy_button=True
            )
            plot_output = gr.Plot(label="Confidence Scores")
    
    with gr.Accordion("📚 Model Details", open=False):
        gr.Markdown("""
        ### **ALL**
        * **Description**: Runs all available detectors (R50_TF, R50_nodown, CLIP-D, P2G, NPR, PGC, ForensicConcept, MDMF) sequentially on the input image. The complete code for training and testing these models on your own datasets is available in our [GitHub repository](https://github.com/truebees-ai/Image-Deepfake-Detectors-Public-Library).
        * **Results**: Shows the prediction of each detector and a **Confidence Plot** with each model's score (no aggregated verdict).

        ### **R50_TF**
        * **Description**: A ResNet50 architecture modified to exclude downsampling at the first layer. It uses "learned prototypes" in the classification head for robust detection.
        * **Paper**: [TrueFake: A Real World Case Dataset of Last Generation Fake Images also Shared on Social Networks](https://arxiv.org/pdf/2504.20658)
        * **Code**: [GitHub Repository](https://github.com/MMLab-unitn/TrueFake-IJCNN25)

        ### **R50_nodown**
        * **Description**: A ResNet-50 model without downsampling operations in the first layer, designed to preserve high-frequency artifacts common in synthetic images.
        * **Paper**: [On the detection of synthetic images generated by diffusion models](https://arxiv.org/abs/2211.00680)
        * **Code**: [GitHub Repository](https://grip-unina.github.io/DMimageDetection/)

        ### **CLIP-D**
        * **Description**: A lightweight detection strategy based on CLIP features. It exhibits surprising generalization ability using only a handful of example images.
        * **Paper**: [Raising the Bar of AI-generated Image Detection with CLIP](https://arxiv.org/abs/2312.00195v2)
        * **Code**: [GitHub Repository](https://grip-unina.github.io/ClipBased-SyntheticImageDetection/)

        ### **P2G (Prompt2Guard)**
        * **Description**: Uses Vision-Language Models (VLMs) with conditioned prompt-optimization for continual deepfake detection. It leverages read-only prompts for efficiency.
        * **Paper**: [Conditioned Prompt-Optimization for Continual Deepfake Detection](https://arxiv.org/abs/2407.21554)
        * **Code**: [GitHub Repository](https://github.com/laitifranz/Prompt2Guard)

        ### **NPR**
        * **Description**: Focuses on Neighboring Pixel Relationships (NPR) to capture generalized structural artifacts stemming from up-sampling operations in generative networks.
        * **Paper**: [Rethinking the Up-Sampling Operations in CNN-based Generative Network for Generalizable Deepfake Detection](https://arxiv.org/abs/2312.10461)
        * **Code**: [GitHub Repository](https://github.com/chuangchuangtan/NPR-DeepfakeDetection)

        ### **PGC**
        * **Description**: Peak-guided calibration on a DINOv2-L backbone. Uses the **official pre-trained weights** released by the authors (not trained on DeepShield).
        * **Weights**: [xiaoyuzhou68/PGC_ckpt](https://huggingface.co/xiaoyuzhou68/PGC_ckpt)

        ### **ForensicConcept**
        * **Description**: CLIP ViT-L/14 with a forensic concept codebook. Uses the **official pre-trained weights** released by the authors (not trained on DeepShield).
        * **Weights**: [ethan225/ForensicConcept](https://huggingface.co/ethan225/ForensicConcept) · **Code**: [GitHub Repository](https://github.com/EthanAdamm/FORENSICCONCEPT)

        ### **MDMF**
        * **Description**: Local distribution shift measured with MMD on DINOv2-L patch embeddings. Uses the **official pre-trained weights**; the real-image reference bank and decision threshold were built locally from DeepShield data.
        * **Code**: [GitHub Repository](https://github.com/ZBox1005/MDMF)
        """)

    gr.Markdown("""
    ---
    ### References
    1. Dell'Anna, S., Montibeller, A., & Boato, G. (2025). *TrueFake: A Real World Case Dataset of Last Generation Fake Images also Shared on Social Networks*. arXiv preprint arXiv:2504.20658.
    2. Corvi, R., et al. (2023). *On the detection of synthetic images generated by diffusion models*. ICASSP.
    3. Cozzolino, D., et al. (2023). *Raising the Bar of AI-generated Image Detection with CLIP*. CVPRW.
    4. Laiti, F., et al. (2024). *Conditioned Prompt-Optimization for Continual Deepfake Detection*. arXiv preprint arXiv:2407.21554.
    5. Tan, C., et al. (2024). *Rethinking the up-sampling operations in cnn-based generative network for generalizable deepfake detection*. CVPR.
    """)
    
    submit_btn.click(
        fn=predict,
        inputs=[image_input, detector_input],
        outputs=[output_display, plot_output]
    )

if __name__ == "__main__":
    # For HF Spaces, configure server settings
    if os.environ.get("SPACE_ID"):
        demo.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["."])
    else:
        # Local execution
        demo.launch()