#!/usr/bin/env python3
import importlib.util
import os
import sys
import time
import cv2
import torch
import numpy as np
import gradio as gr
from PIL import Image
from torchvision import transforms
from torchvision.models import vit_b_16
import torch.nn as nn
import traceback

# Add current directory to path
if not os.getcwd() in sys.path:
    sys.path.append(os.getcwd())

# Check if detectron2 is installed
if importlib.util.find_spec("detectron2") is None:
    print("Installing PyTorch and Detectron2...")
    os.system("pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu")
    os.system("pip install git+https://github.com/facebookresearch/detectron2.git")
    print("Installation complete!")

# Check for detectron2
try:
    from detectron2.engine import DefaultPredictor
    from detectron2.config import get_cfg
    from detectron2.utils.visualizer import Visualizer, ColorMode
    from detectron2 import model_zoo
    DETECTRON2_AVAILABLE = True
except ImportError:
    print("Warning: Detectron2 is not installed. Damage detection will not be available.")
    DETECTRON2_AVAILABLE = False

# Define model paths
DEFAULT_DAMAGE_MODEL_PATH = "./model_final.pth"
DEFAULT_DEEPFAKE_MODEL_PATH = "./vit_deepfake_best.pth"

# Sample images for demo (add your own paths)
SAMPLE_IMAGES = [
    "./test3.png",
    "./test5.png",
]

# Maximum number of tries allowed
MAX_TRIES = 3

def verify_detectron2_installation():
    """Verify that Detectron2 is properly installed"""
    results = {
        "detectron2_installed": False,
        "model_zoo_accessible": False,
        "can_create_cfg": False,
        "error_messages": []
    }
    
    try:
        import importlib.util
        if importlib.util.find_spec("detectron2") is not None:
            results["detectron2_installed"] = True
            
            try:
                import detectron2
                from detectron2 import model_zoo
                config_file = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
                config_path = model_zoo.get_config_file(config_file)
                if os.path.exists(config_path):
                    results["model_zoo_accessible"] = True
            except Exception as e:
                results["error_messages"].append(f"Error accessing model zoo: {str(e)}")
            
            try:
                from detectron2.config import get_cfg
                cfg = get_cfg()
                results["can_create_cfg"] = True
            except Exception as e:
                results["error_messages"].append(f"Error creating Detectron2 config: {str(e)}")
        else:
            results["error_messages"].append("Detectron2 is not installed")
    except Exception as e:
        results["error_messages"].append(f"Error checking Detectron2 installation: {str(e)}")
    
    return results

def setup_device(device_str):
    """Set up the computation device"""
    if device_str == 'auto':
        if torch.cuda.is_available():
            return torch.device('cuda:0')
        elif hasattr(torch, 'backends') and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            return torch.device('mps')
        else:
            return torch.device('cpu')
    elif device_str == 'cuda' and torch.cuda.is_available():
        return torch.device('cuda:0')
    elif device_str == 'mps' and hasattr(torch, 'backends') and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        return torch.device('mps')
    else:
        print(f"Warning: Device {device_str} not available, using CPU instead.")
        return torch.device('cpu')

def setup_damage_detector(model_path, threshold=0.7):
    """Set up the damage detection model"""
    if not DETECTRON2_AVAILABLE:
        print("Detectron2 is not installed. Cannot set up damage detector.")
        return None, None
    
    try:
        print(f"Checking model path: {model_path}")
        print(f"Model exists: {os.path.exists(model_path)}")
        
        if model_path is None or not os.path.exists(model_path):
            print(f"Error: Damage model file not found at {model_path}")
            return None, None
        
        cfg = get_cfg()
        cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
        cfg.MODEL.WEIGHTS = model_path
        cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # Only one class (damage)
        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = threshold
        
        # Use CPU if on Mac (MPS)
        cfg.MODEL.DEVICE = "cpu"
        print("Forcing Detectron2 to use CPU")        
        
        predictor = DefaultPredictor(cfg)
        return predictor, cfg
    except Exception as e:
        print(f"Detailed error: {str(e)}")
        import traceback
        traceback.print_exc()
        return None, None

def load_vit_deepfake_model(model_path, device):
    """Load the Vision Transformer (ViT) model for deepfake detection"""
    if model_path is None or not os.path.exists(model_path):
        print(f"Error: ViT deepfake model file not found at {model_path}")
        return None
        
    try:
        # Create ViT model with binary classification head
        model = vit_b_16(weights=None)
        
        # Modify the classifier head for binary classification (real vs fake)
        in_features = model.heads.head.in_features
        model.heads.head = nn.Linear(in_features, 2)
        
        # Load weights
        print(f"Loading ViT deepfake model from: {model_path}")
        checkpoint = torch.load(model_path, map_location='cpu')
        
        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
        elif isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
            model.load_state_dict(checkpoint['state_dict'])
        else:
            model.load_state_dict(checkpoint)
            
        # Move model to device and set to evaluation mode
        model = model.to(device)
        model.eval()
        
        return model
    except Exception as e:
        print(f"Error loading ViT deepfake model: {e}")
        traceback.print_exc()
        return None

def preprocess_for_vit(image, device):
    """Preprocess an image for Vision Transformer deepfake detection"""
    try:
        # Ensure image is RGB
        if len(image.shape) == 3 and image.shape[2] == 3:
            if image.dtype != np.uint8:
                image = (image * 255).astype(np.uint8)
            rgb_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        else:
            rgb_img = image
            
        # Resize to standard ViT input size (224x224)
        img_resized = cv2.resize(rgb_img, (224, 224))
        
        # Apply transforms
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
        img_tensor = transform(Image.fromarray(img_resized)).unsqueeze(0)
        img_tensor = img_tensor.to(device)
        
        return img_tensor
    except Exception as e:
        print(f"Error preprocessing image for ViT: {e}")
        traceback.print_exc()
        return None

def detect_damage(img, damage_detector):
    """Detect damage in an image"""
    try:
        if img is None:
            raise ValueError("Invalid image")
            
        # If no detector, use whole image
        if damage_detector is None:
            h, w = img.shape[:2]
            damage_regions = [{
                "box": (0, 0, w, h),
                "score": 1.0,
                "mask": None
            }]
            return img, None, damage_regions
            
        # Run inference
        outputs = damage_detector(img)
        
        # Get regions
        instances = outputs["instances"].to("cpu")
        boxes = instances.pred_boxes.tensor.numpy() if instances.has("pred_boxes") else []
        scores = instances.scores.numpy() if instances.has("scores") else []
        masks = instances.pred_masks.numpy() if instances.has("pred_masks") else []
        
        damage_regions = []
        for i in range(len(boxes)):
            x1, y1, x2, y2 = map(int, boxes[i])
            damage_regions.append({
                "box": (x1, y1, x2, y2),
                "score": float(scores[i]),
                "mask": masks[i] if len(masks) > i else None
            })
        
        # If no regions found, use whole image
        if not damage_regions:
            h, w = img.shape[:2]
            damage_regions = [{
                "box": (0, 0, w, h),
                "score": 1.0,
                "mask": None
            }]
            
        return img, outputs, damage_regions
    except Exception as e:
        print(f"Error detecting damage: {e}")
        traceback.print_exc()
        
        # Return whole image if error
        if 'img' in locals() and img is not None:
            h, w = img.shape[:2]
            damage_regions = [{
                "box": (0, 0, w, h),
                "score": 1.0,
                "mask": None
            }]
            return img, None, damage_regions
        return None, None, []

def check_deepfake_vit(image, damage_regions, deepfake_model, device, threshold=0.5):
    """Check if damage regions are deepfakes using ViT model"""
    results = []
    
    if deepfake_model is None:
        return []
    
    try:
        # If no damage regions, check entire image
        if not damage_regions:
            img_tensor = preprocess_for_vit(image, device)
            if img_tensor is None:
                return []
                
            # Run inference
            with torch.no_grad():
                outputs = deepfake_model(img_tensor)
                
                # Get predictions
                probabilities = torch.nn.functional.softmax(outputs, dim=1)
                fake_prob = probabilities[0, 1].item()  # Probability of being fake (class 1)
                is_fake = fake_prob > threshold
                
                results.append({
                    "region": "full_image",
                    "deepfake_prob": float(fake_prob),
                    "is_fake": bool(is_fake)
                })
            
            return results
        
        # Process each damage region
        for i, region in enumerate(damage_regions):
            x1, y1, x2, y2 = region["box"]
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(image.shape[1], x2), min(image.shape[0], y2)
            
            # Only process valid regions
            if x2 > x1 and y2 > y1:
                # Extract region
                roi = image[y1:y2, x1:x2]
                
                # Preprocess
                img_tensor = preprocess_for_vit(roi, device)
                if img_tensor is None:
                    continue
                
                # Inference
                with torch.no_grad():
                    outputs = deepfake_model(img_tensor)
                    
                    # Get predictions
                    probabilities = torch.nn.functional.softmax(outputs, dim=1)
                    fake_prob = probabilities[0, 1].item()  # Probability of being fake (class 1)
                    is_fake = fake_prob > threshold
                    
                    results.append({
                        "region_id": i,
                        "box": (x1, y1, x2, y2),
                        "deepfake_prob": float(fake_prob),
                        "is_fake": bool(is_fake)
                    })
        
        return results
    except Exception as e:
        print(f"Error in deepfake detection: {e}")
        traceback.print_exc()
        return []

def visualize_results(image, damage_outputs, deepfake_results, damage_threshold):
    """Create visualization of results"""
    try:
        img_copy = image.copy()
        
        # Draw damage detection
        if damage_outputs is not None and DETECTRON2_AVAILABLE:
            try:
                v = Visualizer(img_copy[:, :, ::-1], scale=1.0, instance_mode=ColorMode.IMAGE_BW)
                v = v.draw_instance_predictions(damage_outputs["instances"].to("cpu"))
                result_img = v.get_image()[:, :, ::-1]
                result_img = np.array(result_img, dtype=np.uint8)
            except Exception as e:
                print(f"Error visualizing damage: {e}")
                result_img = img_copy
        else:
            result_img = img_copy
        
        # Add deepfake results
        for result in deepfake_results:
            try:
                if "box" in result:
                    x1, y1, x2, y2 = result["box"]
                    fake_prob = result["deepfake_prob"]
                    is_fake = result["is_fake"]
                    region_id = result.get("region_id", 0)
                    
                    # Status text
                    text = f"R{region_id}: {'FAKE' if is_fake else 'REAL'} ({fake_prob*100:.1f}%)"
                    
                    # Red for fake, green for real
                    color = (0, 0, 255) if is_fake else (0, 255, 0)
                    
                    # Ensure standard numpy array
                    if not isinstance(result_img, np.ndarray):
                        result_img = np.array(result_img, dtype=np.uint8)
                    
                    # Draw rectangle and text
                    cv2.rectangle(result_img, (x1, y1), (x2, y2), color, 2)
                    cv2.putText(result_img, text, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
                elif "region" in result and result["region"] == "full_image":
                    fake_prob = result["deepfake_prob"]
                    is_fake = result["is_fake"]
                    
                    text = f"Image: {'FAKE' if is_fake else 'REAL'} ({fake_prob*100:.1f}%)"
                    color = (0, 0, 255) if is_fake else (0, 255, 0)
                    
                    if not isinstance(result_img, np.ndarray):
                        result_img = np.array(result_img, dtype=np.uint8)
                    
                    cv2.putText(result_img, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
            except Exception as e:
                print(f"Error drawing result: {e}")
        
        return result_img
    except Exception as e:
        print(f"Error in visualization: {e}")
        traceback.print_exc()
        return np.array(image, dtype=np.uint8)

def process_image(input_image, damage_model_path, deepfake_model_path, damage_threshold, 
                 deepfake_threshold, skip_damage, device_str, usage_count):
    """Process an image through the detection pipeline"""
    # Increment usage count and check if limit reached
    usage_count = usage_count + 1
    
    progress_info = []
    progress_info.append(f"Usage: {usage_count}/{MAX_TRIES}")
    
    # Check model files
    if not skip_damage and damage_model_path:
        if not os.path.exists(damage_model_path):
            progress_info.append(f"ERROR: Damage model not found at {damage_model_path}")
    
    if deepfake_model_path and not os.path.exists(deepfake_model_path):
        progress_info.append(f"ERROR: ViT deepfake model not found at {deepfake_model_path}")
    
    # Convert image to proper format
    try:
        if isinstance(input_image, dict) and "path" in input_image:
            img = cv2.imread(input_image["path"])
        elif isinstance(input_image, str):
            img = cv2.imread(input_image)
        elif isinstance(input_image, np.ndarray):
            img = input_image.copy()
            if len(img.shape) == 3 and img.shape[2] == 3:
                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        else:
            return None, "Error: Unsupported image format", usage_count
        
        if img is None:
            return None, "Error: Could not read the image", usage_count
    except Exception as e:
        return None, f"Error loading image: {str(e)}", usage_count
    
    # Setup device
    device = setup_device(device_str)
    progress_info.append(f"Using device: {device}")
    
    # Initialize models
    damage_detector = None
    deepfake_model = None
    
    # Setup damage detector
    if not skip_damage and damage_model_path:
        progress_info.append("Setting up damage detector...")
        damage_detector, _ = setup_damage_detector(damage_model_path, float(damage_threshold))
        if damage_detector:
            progress_info.append("✅ Damage detector initialized")
        else:
            progress_info.append("❌ Failed to initialize damage detector")
        
    # Setup ViT deepfake detector
    if deepfake_model_path:
        progress_info.append("Setting up ViT deepfake detector...")
        deepfake_model = load_vit_deepfake_model(deepfake_model_path, device)
        if deepfake_model:
            progress_info.append("✅ ViT deepfake detector initialized")
        else:
            progress_info.append("❌ Failed to initialize ViT deepfake detector")
     
    # Step 1: Detect damage
    progress_info.append("Detecting damaged regions...")
    start_time = time.time()
    img, damage_outputs, damage_regions = detect_damage(img, damage_detector)
    damage_time = time.time() - start_time
    
    if damage_regions:
        progress_info.append(f"Found {len(damage_regions)} damage regions in {damage_time:.2f} seconds")
    else:
        progress_info.append("No damage regions detected")
    
    # Step 2: Check for deepfakes using ViT
    deepfake_results = []
    if deepfake_model is not None:
        progress_info.append("Analyzing regions for deepfakes using ViT model...")
        start_time = time.time()
        deepfake_results = check_deepfake_vit(
            img, damage_regions, deepfake_model, device, float(deepfake_threshold)
        )
        deepfake_time = time.time() - start_time
        
        if deepfake_results:
            progress_info.append(f"Deepfake analysis completed in {deepfake_time:.2f} seconds")
            
            # Generate report
            for result in deepfake_results:
                if "region_id" in result:
                    region_id = result["region_id"]
                    fake_prob = result["deepfake_prob"]
                    is_fake = result["is_fake"]
                    progress_info.append(f"Region {region_id}: {'FAKE' if is_fake else 'REAL'} (Probability: {fake_prob*100:.2f}%)")
                elif "region" in result and result["region"] == "full_image":
                    fake_prob = result["deepfake_prob"]
                    is_fake = result["is_fake"]
                    progress_info.append(f"Whole image: {'FAKE' if is_fake else 'REAL'} (Probability: {fake_prob*100:.2f}%)")
        else:
            progress_info.append("No deepfake detection results")
    
    # Final verdict
    fake_regions = [r for r in deepfake_results if r.get("is_fake", False)]
    if fake_regions:
        progress_info.append("\n🚨 VERDICT: This image contains FAKE damage 🚨")
    else:
        progress_info.append("\n✅ VERDICT: All damage appears REAL")
        
    # Step 3: Visualize results
    result_img = visualize_results(img, damage_outputs, deepfake_results, float(damage_threshold))
    
    # Convert back to RGB for Gradio
    if len(result_img.shape) == 3 and result_img.shape[2] == 3:
        result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
    
    # Add usage information at the bottom of the image
    if usage_count >= MAX_TRIES:
        # Add a "Usage limit reached" message to the bottom of the image in red
        cv2.putText(result_img, f"USAGE LIMIT REACHED: {usage_count}/{MAX_TRIES}", 
                    (10, result_img.shape[0] - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    else:
        # Add a usage counter to the bottom of the image
        cv2.putText(result_img, f"Usage: {usage_count}/{MAX_TRIES}", 
                    (10, result_img.shape[0] - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
    
    # Add usage info to the progress text
    if usage_count >= MAX_TRIES:
        progress_info.append("\n⚠️ You have reached the maximum number of tries allowed ⚠️")
    else:
        progress_info.append(f"\nRemaining tries: {MAX_TRIES - usage_count}")
    
    return result_img, "\n".join(progress_info), usage_count

def auto_install_dependencies():
    """Attempt to install dependencies if needed"""
    try:
        import importlib.util
        
        # Check for PyTorch
        if importlib.util.find_spec("torch") is None:
            print("Installing PyTorch...")
            os.system("pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu")
        
        # Check for Detectron2
        if importlib.util.find_spec("detectron2") is None:
            print("Installing Detectron2...")
            os.system("pip install git+https://github.com/facebookresearch/detectron2.git")
        
        # Check for Gradio
        if importlib.util.find_spec("gradio") is None:
            print("Installing Gradio...")
            os.system("pip install gradio")
        
        print("Dependencies installation complete!")
        return True
    except Exception as e:
        print(f"Error installing dependencies: {e}")
        return False

def check_model_paths(damage_path, deepfake_path):
    """Check if model paths are valid and exist"""
    output = ["## Path Verification Results\n"]
    
    # Check damage model
    if os.path.exists(damage_path):
        file_size = os.path.getsize(damage_path) / (1024 * 1024)  # Size in MB
        output.append(f"✅ **Damage model:** Found at {damage_path} ({file_size:.2f} MB)")
    else:
        output.append(f"❌ **Damage model:** NOT found at {damage_path}")
    
    # Check deepfake model
    if os.path.exists(deepfake_path):
        file_size = os.path.getsize(deepfake_path) / (1024 * 1024)  # Size in MB
        output.append(f"✅ **ViT deepfake model:** Found at {deepfake_path} ({file_size:.2f} MB)")
    else:
        output.append(f"❌ **ViT deepfake model:** NOT found at {deepfake_path}")
    
    return "\n".join(output)

def create_gradio_interface():
    # Define a theme
    theme = gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="orange",
    )
    
    # Initialize usage counter
    usage_counter = gr.State(0)
    
    with gr.Blocks(title="Car Damage & Deepfake Detector", theme=theme) as app:
        gr.Markdown("""
        # 🚗 Car Damage Fraud Detector with Vision Transformer
        
        Upload a car image to:
        1. Detect damaged areas
        2. Verify if the damage is real or artificially generated (deepfake)
        
        *This app uses a Vision Transformer (ViT) model for deepfake detection.*
        
        ⚠️ **Note: You have a maximum of 3 tries to analyze images.**
        """)
        
        # Main Interface Tab
        with gr.Tab("Analyze Image"):
            with gr.Row():
                with gr.Column(scale=1):
                    input_image = gr.Image(type="numpy", label="Upload Car Image")
                    
                    with gr.Row():
                        process_btn = gr.Button("Analyze Image", variant="primary")
                        clear_btn = gr.Button("Clear", variant="secondary")
                    
                    # Usage limit display
                    usage_display = gr.Markdown("**Usage: 0/3**")
                    
                    with gr.Accordion("Advanced Settings", open=False):
                        skip_damage = gr.Checkbox(label="Skip Damage Detection", value=False)
                        damage_threshold = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.05,
                                                   label="Damage Detection Threshold")
                        deepfake_threshold = gr.Slider(minimum=0.1, maximum=1.0, value=0.5, step=0.05,
                                                     label="Deepfake Detection Threshold")
                        device = gr.Dropdown(choices=["auto", "cuda", "cpu", "mps"], value="auto",
                                            label="Computation Device")
                        
                with gr.Column(scale=1):
                    output_image = gr.Image(type="numpy", label="Analysis Result")
                    
                    # Analysis info with nice formatting
                    with gr.Accordion("Analysis Details", open=True):
                        output_text = gr.Markdown(label="Detection Results")
        
 # This is the continuation of the code, starting from the System Diagnostics tab

        with gr.Tab("System Diagnostics"):
            with gr.Row():
                run_diagnostic_btn = gr.Button("Run Diagnostics", variant="primary")
                install_deps_btn = gr.Button("Install Dependencies", variant="secondary")
                
            diagnostic_output = gr.Markdown(label="Diagnostic Results")
            
            # Function to run system diagnostics
            def run_diagnostics():
                detectron2_results = verify_detectron2_installation()
                
                output = ["## System Diagnostics\n"]
                
                # Python & PyTorch versions
                import sys
                output.append(f"**Python:** {sys.version.split()[0]}")
                
                try:
                    import torch
                    output.append(f"**PyTorch:** {torch.__version__}")
                    output.append(f"**CUDA Available:** {'Yes ✅' if torch.cuda.is_available() else 'No ❌'}")
                    if torch.cuda.is_available():
                        output.append(f"**GPU:** {torch.cuda.get_device_name(0)}")
                    output.append(f"**MPS Available:** {'Yes ✅' if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available() else 'No ❌'}")
                except ImportError:
                    output.append("**PyTorch:** Not installed ❌")
                
                output.append("\n## Model Files")
                
                # Check model files
                if os.path.exists(DEFAULT_DAMAGE_MODEL_PATH):
                    file_size = os.path.getsize(DEFAULT_DAMAGE_MODEL_PATH) / (1024 * 1024)  # Size in MB
                    output.append(f"**Damage Model:** Available ✅ ({file_size:.2f} MB)")
                else:
                    output.append(f"**Damage Model:** Not found ❌ at {DEFAULT_DAMAGE_MODEL_PATH}")
                
                if os.path.exists(DEFAULT_DEEPFAKE_MODEL_PATH):
                    file_size = os.path.getsize(DEFAULT_DEEPFAKE_MODEL_PATH) / (1024 * 1024)  # Size in MB
                    output.append(f"**ViT Deepfake Model:** Available ✅ ({file_size:.2f} MB)")
                else:
                    output.append(f"**ViT Deepfake Model:** Not found ❌ at {DEFAULT_DEEPFAKE_MODEL_PATH}")
                
                output.append("\n## Detectron2 Status")
                output.append(f"**Installed:** {'Yes ✅' if detectron2_results['detectron2_installed'] else 'No ❌'}")
                output.append(f"**Model Zoo Access:** {'Yes ✅' if detectron2_results['model_zoo_accessible'] else 'No ❌'}")
                output.append(f"**Config Creation:** {'Yes ✅' if detectron2_results['can_create_cfg'] else 'No ❌'}")
                
                if detectron2_results['error_messages']:
                    output.append("\n## Error Messages")
                    for error in detectron2_results['error_messages']:
                        output.append(f"- {error}")
                
                output.append("\n## Recommendations")
                recommendations = []
                
                if not detectron2_results['detectron2_installed']:
                    recommendations.append("Install Detectron2: `pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu`")
                    recommendations.append("Then: `pip install git+https://github.com/facebookresearch/detectron2.git`")
                
                if not os.path.exists(DEFAULT_DAMAGE_MODEL_PATH):
                    recommendations.append(f"Place the damage model file at: {DEFAULT_DAMAGE_MODEL_PATH}")
                
                if not os.path.exists(DEFAULT_DEEPFAKE_MODEL_PATH):
                    recommendations.append(f"Place the ViT deepfake model file at: {DEFAULT_DEEPFAKE_MODEL_PATH}")
                
                if recommendations:
                    for rec in recommendations:
                        output.append(f"- {rec}")
                else:
                    output.append("- All systems are ready! 🎉")
                
                return "\n".join(output)
            
            # Function to install dependencies
            def install_dependencies():
                output = ["## Installing Dependencies\n"]
                
                # Try to install PyTorch
                output.append("Installing PyTorch...")
                try:
                    import importlib.util
                    if importlib.util.find_spec("torch") is None:
                        os.system("pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu")
                        output.append("✅ PyTorch installed")
                    else:
                        output.append("✅ PyTorch already installed")
                except Exception as e:
                    output.append(f"❌ Error installing PyTorch: {str(e)}")
                
                # Try to install Detectron2
                output.append("\nInstalling Detectron2...")
                try:
                    if importlib.util.find_spec("detectron2") is None:
                        os.system("pip install git+https://github.com/facebookresearch/detectron2.git")
                        output.append("✅ Detectron2 installed")
                    else:
                        output.append("✅ Detectron2 already installed")
                except Exception as e:
                    output.append(f"❌ Error installing Detectron2: {str(e)}")
                
                output.append("\n**Note:** You may need to restart the application for changes to take effect.")
                
                return "\n".join(output)
        
        # Settings Tab
        with gr.Tab("Settings"):
            gr.Markdown("## Model Settings")
            
            with gr.Row():
                with gr.Column():
                    custom_damage_model = gr.Textbox(
                        label="Custom Damage Model Path",
                        value=DEFAULT_DAMAGE_MODEL_PATH,
                        placeholder="Path to damage detection model (.pth)"
                    )
                    
                    custom_deepfake_model = gr.Textbox(
                        label="Custom ViT Deepfake Model Path",
                        value=DEFAULT_DEEPFAKE_MODEL_PATH,
                        placeholder="Path to ViT deepfake detection model (.pth)"
                    )
                    
                    check_paths_btn = gr.Button("Verify Paths", variant="primary")
                
                with gr.Column():
                    paths_result = gr.Markdown(label="Path Verification Results")

        with gr.Tab("Help"):
            gr.Markdown("""
            ## 📋 How to Use This Tool
            
            ### Basic Usage
            1. Upload a car image in the "Analyze Image" tab
            2. Click "Analyze Image" to process it
            3. View the results - damaged areas will be highlighted and classified as real or fake
            
            ### Understanding Results
            - **Green boxes/text:** Real damage
            - **Red boxes/text:** Potential deepfake damage
            - Percentage values show the confidence score
            
            ### Requirements
            - Damage detection model file (Detectron2 Mask R-CNN)
            - Vision Transformer (ViT) deepfake detection model file
            
            ### About the Vision Transformer (ViT)
            - This version uses a state-of-the-art Vision Transformer model for deepfake detection
            - The ViT model has been trained specifically to detect artificially generated car damage
            - ViT models are better at capturing global features in images compared to traditional CNNs
            
            ### Troubleshooting
            - If models aren't loading, check the "System Diagnostics" tab
            - Ensure all model files are in the expected locations
            - Try installing dependencies from the Diagnostics tab
            
            ### Advanced Settings
            - **Damage Threshold:** Higher values mean only high-confidence damage regions are shown
            - **Deepfake Threshold:** Higher values make the system more selective in flagging fakes
            - **Skip Damage Detection:** Analyze the entire image for deepfakes without damage detection
            """)
            
        # Examples
        if any(os.path.exists(img) for img in SAMPLE_IMAGES):
            gr.Markdown("## Example Images")
            with gr.Row():
                example_inputs = [img for img in SAMPLE_IMAGES if os.path.exists(img)]
                gr.Examples(
                    examples=example_inputs,
                    inputs=input_image,
                    outputs=[output_image, output_text, usage_counter],  # Add usage_counter to outputs
                    fn=lambda x: process_image(
                        x, DEFAULT_DAMAGE_MODEL_PATH, DEFAULT_DEEPFAKE_MODEL_PATH, 
                        0.7, 0.5, False, "auto", 0  # Added initial usage_count value of 0
                    ),
                    cache_examples=True
                )
        
        # Connect functions to the UI
        process_btn.click(
            fn=process_image,
            inputs=[
                input_image, 
                custom_damage_model,
                custom_deepfake_model,
                damage_threshold,
                deepfake_threshold,
                skip_damage,
                device,
                usage_counter  # Add the usage counter state here
            ],
            outputs=[output_image, output_text, usage_counter]  # Update usage counter in outputs
        )
        
        # Clear button functionality
        clear_btn.click(
            fn=lambda: [None, "", 0],  # Reset usage counter to 0 when clearing
            inputs=[],
            outputs=[output_image, output_text, usage_counter]
        )
        
        # System diagnostic buttons
        run_diagnostic_btn.click(
            fn=run_diagnostics,
            inputs=[],
            outputs=diagnostic_output
        )
        
        install_deps_btn.click(
            fn=install_dependencies,
            inputs=[],
            outputs=diagnostic_output
        )
        
        # Settings tab
        check_paths_btn.click(
            fn=check_model_paths,
            inputs=[custom_damage_model, custom_deepfake_model],
            outputs=paths_result
        )
        
        # Update usage display when counter changes
        usage_counter.change(
            fn=lambda count: f"**Usage: {count}/{MAX_TRIES}**",
            inputs=[usage_counter],
            outputs=[usage_display]
        )
        
    return app

if __name__ == "__main__":
    # Check if dependencies are installed
    auto_install_dependencies()
    
    # Create and launch the Gradio app
    app = create_gradio_interface()
    app.launch(share=False)  # Set share=True to create a public link