# 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']

DETECTOR_WEIGHTS = {
    'CLIP-D': 0.30,
    'R50_TF': 0.25,
    'R50_nodown': 0.20,
    'P2G': 0.15,
    'NPR': 0.10
}

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
        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
            
            votes_real = 0.0
            votes_fake = 0.0
            total_weight_used = 0.0
            confidences = []
            labels = []
            colors = []
            
            for det, res in results:
                pred = res.get('prediction', 'Unknown')
                raw_conf = res.get('confidence', 0.0)
                
                # Calculate 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)
                
                # Weighted Voting logic
                # Only count vote if confidence > 0.6
                if score > 0.6:
                    weight = DETECTOR_WEIGHTS.get(det, 0.0)
                    if pred == 'fake':
                        votes_fake += weight * score
                        total_weight_used += weight
                    elif pred == 'real':
                        votes_real += weight * score
                        total_weight_used += weight
            
            # Majority Voting
            if votes_real > votes_fake:
                verdict = "REAL"
            elif votes_fake > votes_real:
                verdict = "FAKE"
            else:
                verdict = "UNCERTAIN"
            
            # Calculate weighted average confidence
            if total_weight_used > 0:
                weighted_conf = (votes_real + votes_fake) / total_weight_used
            else:
                weighted_conf = 0.0
            
            # Explanation
            if verdict == "REAL":
                explanation = f"Considering the results obtained by all models (weighted by their historical performance), the analyzed image results, with a weighted confidence of {weighted_conf:.4f}, not produced by a generative AI."
            elif verdict == "FAKE":
                explanation = f"Considering the results obtained by all models (weighted by their historical performance), the analyzed image results, with a weighted confidence of {weighted_conf:.4f}, produced by a generative AI."
            else:
                explanation = f"The result is uncertain. The detectors produced unconsistent results. The weighted confidence is {weighted_conf:.4f}."
            
            # 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')
            ax.axhline(y=0.6, color='gray', linestyle='--', alpha=0.5, label='Vote Threshold (0.6)')
            
            # Add custom legend for colors
            from matplotlib.patches import Patch
            legend_elements = [
                Patch(facecolor='green', label='Real'),
                Patch(facecolor='red', label='Fake'),
                ax.lines[0] # The threshold line
            ]
            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('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")
    gr.Markdown("""
    This space collects a series of state-of-the-art methods for deepfake detection, allowing for free and unlimited use.
    
    ### Training & Performance
    All methods have been trained using the **[DeepShield dataset](https://zenodo.org/records/15648378)**, on images generated with **Stable Diffusion XL** and **StyleGAN 2**. 
    You can expect performance comparable to the results shown in [Dell'Anna et al. (2025)](https://arxiv.org/pdf/2504.20658).

    ### Understanding the Results
    *   **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).

    ### Understanding the Results produced by "ALL"
    *   Runs all available detectors (R50_TF, R50_nodown, CLIP-D, P2G, NPR) sequentially on the input image.
    *   Produces a **Weighted Majority Vote** verdict (Real/Fake). Each model's vote is weighted by a fixed importance score (summing to 1) based on user ranking **and its confidence score**. Only confident predictions (> 0.6) are counted.
    *   You can find the specific weights used for each model in the **"⚖️ Weight Details"** menu below.
    *   Also generates a **Confidence Plot** visualizing each model's score and a textual **Explanation** of the consensus.
    *   In the plot, **Green** bars indicate a **Real** prediction, while **Red** bars indicate a **Fake** prediction.


    ### Note
    ⚠️ 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.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("⚖️ Weight Details", open=False):
        gr.Markdown(f"""
        ### **Detector Weights**
        The weights are assigned based on the ranking (based on the results of [TrueFake: A Real World Case Dataset of Last Generation Fake Images also Shared on Social Networks](https://arxiv.org/pdf/2504.20658)): **CLIP-D > R50_TF > R50_nodown > P2G > NPR**, such that their sum equals 1.
        
        | Detector | Weight |
        | :--- | :---: |
        | **CLIP-D** | {DETECTOR_WEIGHTS['CLIP-D']:.2f} |
        | **R50_TF** | {DETECTOR_WEIGHTS['R50_TF']:.2f} |
        | **R50_nodown** | {DETECTOR_WEIGHTS['R50_nodown']:.2f} |
        | **P2G** | {DETECTOR_WEIGHTS['P2G']:.2f} |
        | **NPR** | {DETECTOR_WEIGHTS['NPR']:.2f} |
        """)

    with gr.Accordion("📚 Model Details", open=False):
        gr.Markdown("""
        ### **ALL**
        *   **Description**: Runs all available detectors (R50_TF, R50_nodown, CLIP-D, P2G, NPR) sequentially on the input image.
        *   **Results**: Produces a **Majority Vote** verdict (Real/Fake) considering only confident predictions (> 0.6). Also generates a **Confidence Plot** visualizing each model's score and a textual **Explanation** of the consensus.

        ### **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)
        """)

    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()