# =============================================================================
# FACE CLASSIFIER TESTING PROGRAM
# Tests trained model on images with face detection and cropping
# =============================================================================

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import os
import matplotlib.pyplot as plt
from pathlib import Path
import time
from tqdm import tqdm

# =============================================================================
# CONFIGURATION
# =============================================================================

# Paths
MODEL_PATH = r"../Training/best_face_classifier_real_data.pth"
TEST_IMAGES_PATH = r"\Pictures\Saved Pictures"
OUTPUT_PATH = "test_results"

# Model parameters (must match training configuration)
IMAGE_SIZE = 224
INPUT_CHANNELS = 3
NUM_CLASSES = 1
CONV_FILTERS = [128, 256, 512]  # Updated to match TrainV3.py
FC_SIZES = [1024, 512]
DROPOUT_RATES = [0.3, 0.5]

# Image processing
FACE_DETECTION_SCALE_FACTOR = 1.1
FACE_DETECTION_MIN_NEIGHBORS = 5
MIN_FACE_SIZE = (30, 30)
IMAGE_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']

# Visualization
CONFIDENCE_THRESHOLD = 0.8
SAVE_RESULTS = True
SHOW_PLOTS = True
SAVE_INDIVIDUAL_IMAGES = True  # Save each image with annotations
CREATE_COMPREHENSIVE_SUMMARY = True  # Create complete grid summary

# =============================================================================
# MODEL ARCHITECTURE (Must match training)
# =============================================================================

class ImprovedFaceClassifierCNN(nn.Module):
    """Same architecture as used in training"""

    def __init__(self):
        super().__init__()

        # Feature extraction layers
        self.features = nn.Sequential(
            # Block 1: 224x224 -> 112x112
            nn.Conv2d(INPUT_CHANNELS, CONV_FILTERS[0], 3, padding=1),
            nn.BatchNorm2d(CONV_FILTERS[0]),
            nn.ReLU(inplace=True),
            nn.Conv2d(CONV_FILTERS[0], CONV_FILTERS[0], 3, padding=1),
            nn.BatchNorm2d(CONV_FILTERS[0]),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
            nn.Dropout(DROPOUT_RATES[0]),

            # Block 2: 112x112 -> 56x56
            nn.Conv2d(CONV_FILTERS[0], CONV_FILTERS[1], 3, padding=1),
            nn.BatchNorm2d(CONV_FILTERS[1]),
            nn.ReLU(inplace=True),
            nn.Conv2d(CONV_FILTERS[1], CONV_FILTERS[1], 3, padding=1),
            nn.BatchNorm2d(CONV_FILTERS[1]),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
            nn.Dropout(DROPOUT_RATES[0]),

            # Block 3: 56x56 -> 28x28
            nn.Conv2d(CONV_FILTERS[1], CONV_FILTERS[2], 3, padding=1),
            nn.BatchNorm2d(CONV_FILTERS[2]),
            nn.ReLU(inplace=True),
            nn.Conv2d(CONV_FILTERS[2], CONV_FILTERS[2], 3, padding=1),
            nn.BatchNorm2d(CONV_FILTERS[2]),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
            nn.Dropout(DROPOUT_RATES[0]),

            # Global Average Pooling
            nn.AdaptiveAvgPool2d((7, 7))
        )

        # Classifier
        self.classifier = nn.Sequential(
            nn.Linear(CONV_FILTERS[2] * 7 * 7, FC_SIZES[0]),
            nn.BatchNorm1d(FC_SIZES[0]),
            nn.ReLU(inplace=True),
            nn.Dropout(DROPOUT_RATES[1]),

            nn.Linear(FC_SIZES[0], FC_SIZES[1]),
            nn.ReLU(inplace=True),
            nn.Dropout(DROPOUT_RATES[1]),

            nn.Linear(FC_SIZES[1], NUM_CLASSES)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

# =============================================================================
# FACE DETECTION AND PROCESSING
# =============================================================================

class FaceProcessor:
    """Face detection and processing for classification"""

    def __init__(self):
        # Initialize face detector (OpenCV Haar Cascade)
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

        # Alternative: Try to use more accurate DNN face detector if available
        try:
            # Download OpenCV DNN face detector if not present
            self.net = None
            self.use_dnn = False
            # Note: For production, you might want to use a more sophisticated face detector
        except:
            self.net = None
            self.use_dnn = False

        # Image preprocessing transform
        self.transform = transforms.Compose([
            transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def detect_faces(self, image):
        """Detect faces in image and return bounding boxes with duplicate removal"""
        if isinstance(image, Image.Image):
            # Convert PIL to OpenCV format
            image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        else:
            image_cv = image.copy()

        # Convert to grayscale for face detection
        gray = cv2.cvtColor(image_cv, cv2.COLOR_BGR2GRAY)

        # Detect faces
        faces = self.face_cascade.detectMultiScale(
            gray,
            scaleFactor=FACE_DETECTION_SCALE_FACTOR,
            minNeighbors=FACE_DETECTION_MIN_NEIGHBORS,
            minSize=MIN_FACE_SIZE,
            flags=cv2.CASCADE_SCALE_IMAGE
        )

        # Remove duplicate/overlapping faces using Non-Maximum Suppression
        if len(faces) > 1:
            faces = self._remove_duplicate_faces(faces)

        return faces

    def _remove_duplicate_faces(self, faces, overlap_threshold=0.15):
        """Remove duplicate/overlapping face detections using improved NMS"""
        if len(faces) <= 1:
            return faces

        # Convert to list for easier manipulation
        face_list = list(faces)
        
        # Calculate areas and create extended info
        face_info = []
        for i, (x, y, w, h) in enumerate(face_list):
            area = w * h
            face_info.append({
                'index': i,
                'bbox': (x, y, w, h),
                'area': area,
                'x1': x, 'y1': y, 'x2': x + w, 'y2': y + h
            })
        
        # Sort by area (larger faces first - usually more reliable)
        face_info.sort(key=lambda f: f['area'], reverse=True)
        
        keep_indices = []
        
        for i, current_face in enumerate(face_info):
            should_keep = True
            
            # Check against all previously kept faces
            for kept_idx in keep_indices:
                kept_face = face_info[kept_idx]
                
                # Calculate intersection
                x1 = max(current_face['x1'], kept_face['x1'])
                y1 = max(current_face['y1'], kept_face['y1'])
                x2 = min(current_face['x2'], kept_face['x2'])
                y2 = min(current_face['y2'], kept_face['y2'])
                
                if x1 < x2 and y1 < y2:
                    intersection = (x2 - x1) * (y2 - y1)
                    
                    # Calculate IoU
                    union = current_face['area'] + kept_face['area'] - intersection
                    iou = intersection / union if union > 0 else 0
                    
                    # Also check overlap ratio (intersection over smaller box)
                    smaller_area = min(current_face['area'], kept_face['area'])
                    overlap_ratio = intersection / smaller_area if smaller_area > 0 else 0
                    
                    # Remove if either IoU or overlap ratio is too high
                    if iou > overlap_threshold or overlap_ratio > 0.5:
                        should_keep = False
                        break
            
            if should_keep:
                keep_indices.append(i)
        
        # Return filtered faces
        filtered_faces = np.array([face_info[i]['bbox'] for i in keep_indices])
        
        # Debug info
        if len(faces) != len(filtered_faces):
            print(f"   [NMS] Removed {len(faces) - len(filtered_faces)} duplicate faces "
                  f"({len(faces)} → {len(filtered_faces)})")
        
        return filtered_faces

    def crop_face(self, image, face_bbox, expand_ratio=0.2):
        """Crop face from image with some padding"""
        x, y, w, h = face_bbox

        # Add padding around face
        pad_x = int(w * expand_ratio)
        pad_y = int(h * expand_ratio)

        # Calculate expanded bounding box
        x1 = max(0, x - pad_x)
        y1 = max(0, y - pad_y)
        x2 = min(image.width if isinstance(image, Image.Image) else image.shape[1], x + w + pad_x)
        y2 = min(image.height if isinstance(image, Image.Image) else image.shape[0], y + h + pad_y)

        # Crop the face
        if isinstance(image, Image.Image):
            face_crop = image.crop((x1, y1, x2, y2))
        else:
            # OpenCV format
            face_crop = image[y1:y2, x1:x2]
            face_crop = Image.fromarray(cv2.cvtColor(face_crop, cv2.COLOR_BGR2RGB))

        return face_crop, (x1, y1, x2, y2)

    def preprocess_face(self, face_image):
        """Preprocess face for model input"""
        # Ensure face is PIL Image
        if not isinstance(face_image, Image.Image):
            face_image = Image.fromarray(face_image)

        # Apply transforms
        face_tensor = self.transform(face_image)

        # Add batch dimension
        face_batch = face_tensor.unsqueeze(0)

        return face_batch

# =============================================================================
# MODEL LOADER AND CLASSIFIER
# =============================================================================

class FaceClassifierTester:
    """Test trained face classifier on new images"""

    def __init__(self, model_path, device='auto'):
        self.device = self._setup_device(device)
        self.model = self._load_model(model_path)
        self.face_processor = FaceProcessor()
        self.results = []

        print(f"[*] Face Classifier Tester initialized")
        print(f"   Device: {self.device}")
        print(f"   Model: {model_path}")

    def _setup_device(self, device):
        """Setup computing device"""
        if device == 'auto':
            if torch.cuda.is_available():
                device = torch.device('cuda:0')
                print(f"[GPU] Using GPU: {torch.cuda.get_device_name(0)}")
            else:
                device = torch.device('cpu')
                print("[CPU] Using CPU")
        else:
            device = torch.device(device)

        return device

    def _load_model(self, model_path):
        """Load trained model from checkpoint"""
        try:
            # Load checkpoint
            checkpoint = torch.load(model_path, map_location=self.device)

            # Initialize model
            model = ImprovedFaceClassifierCNN()

            # Load state dict
            if 'model_state_dict' in checkpoint:
                model.load_state_dict(checkpoint['model_state_dict'])
                print(f"[OK] Model loaded from checkpoint")
                print(f"   Epoch: {checkpoint.get('epoch', 'Unknown')}")
                print(f"   Validation Accuracy: {checkpoint.get('val_acc', 'Unknown'):.2f}%")
            else:
                # Direct state dict
                model.load_state_dict(checkpoint)
                print(f"[OK] Model loaded successfully")

            model.to(self.device)
            model.eval()

            return model

        except Exception as e:
            print(f"[ERROR] Error loading model: {e}")
            print("Make sure the model file exists and matches the architecture")
            raise

    def classify_face(self, face_image):
        """Classify a single face image"""
        try:
            # Preprocess face
            face_tensor = self.face_processor.preprocess_face(face_image)
            face_tensor = face_tensor.to(self.device)

            # Run inference
            with torch.no_grad():
                logits = self.model(face_tensor)
                probability = torch.sigmoid(logits).cpu().numpy()[0][0]

                # Convert probability to prediction
                prediction = "REAL" if probability > CONFIDENCE_THRESHOLD else "FAKE"
                confidence = probability if prediction == "REAL" else (1 - probability)

                return {
                    'prediction': prediction,
                    'confidence': confidence,
                    'probability': probability,
                    'raw_logit': logits.cpu().numpy()[0][0]
                }

        except Exception as e:
            print(f"[ERROR] Error in classification: {e}")
            return {
                'prediction': 'ERROR',
                'confidence': 0.0,
                'probability': 0.0,
                'raw_logit': 0.0
            }

    def process_image(self, image_path):
        """Process a single image: detect faces and classify them"""
        try:
            # Load image
            image = Image.open(image_path).convert('RGB')
            image_name = os.path.basename(image_path)

            print(f"\n[PROCESSING] {image_name}")

            # Detect faces
            faces = self.face_processor.detect_faces(image)

            if len(faces) == 0:
                print(f"   [WARNING] No faces detected in {image_name}")
                return {
                    'image_path': image_path,
                    'image_name': image_name,
                    'num_faces': 0,
                    'faces': [],
                    'status': 'no_faces'
                }

            print(f"   [FACES] Found {len(faces)} face(s)")

            # Process each detected face
            face_results = []
            for i, face_bbox in enumerate(faces):
                # Crop face
                face_crop, expanded_bbox = self.face_processor.crop_face(image, face_bbox)

                # Classify face
                classification = self.classify_face(face_crop)

                # Store results
                face_result = {
                    'face_id': i,
                    'bbox': face_bbox.tolist(),
                    'expanded_bbox': expanded_bbox,
                    'face_crop': face_crop,
                    'classification': classification
                }
                face_results.append(face_result)

                print(f"   Face {i+1}: {classification['prediction']} "
                      f"({classification['confidence']:.1%} confidence)")

            return {
                'image_path': image_path,
                'image_name': image_name,
                'image': image,
                'num_faces': len(faces),
                'faces': face_results,
                'status': 'success'
            }

        except Exception as e:
            print(f"[ERROR] Error processing {image_path}: {e}")
            return {
                'image_path': image_path,
                'image_name': os.path.basename(image_path),
                'num_faces': 0,
                'faces': [],
                'status': 'error',
                'error': str(e)
            }

    def test_folder(self, folder_path, max_images=None):
        """Test all images in a folder"""
        print(f"\n[TESTING] FACE CLASSIFIER")
        print(f"="*60)
        print(f"Test folder: {folder_path}")
        print(f"Model: {MODEL_PATH}")

        # Check if folder exists
        if not os.path.exists(folder_path):
            print(f"[ERROR] Test folder not found: {folder_path}")
            return []

        # Get all image files (avoid duplicates from case variations)
        image_files_set = set()
        for ext in IMAGE_EXTENSIONS:
            # Use case-insensitive glob patterns
            image_files_set.update(Path(folder_path).glob(f"*{ext}"))
            image_files_set.update(Path(folder_path).glob(f"*{ext.upper()}"))
        
        # Convert set back to list and remove duplicates by resolving paths
        image_files = []
        seen_paths = set()
        for file_path in image_files_set:
            resolved_path = file_path.resolve()
            if resolved_path not in seen_paths:
                image_files.append(file_path)
                seen_paths.add(resolved_path)

        if not image_files:
            print(f"[ERROR] No images found in {folder_path}")
            print(f"   Looking for extensions: {IMAGE_EXTENSIONS}")
            return []

        if max_images:
            image_files = image_files[:max_images]

        print(f"[FILES] Found {len(image_files)} images to process")

        # Process each image
        results = []
        start_time = time.time()

        for image_path in tqdm(image_files, desc="Processing images"):
            result = self.process_image(str(image_path))
            results.append(result)
            self.results.append(result)

        total_time = time.time() - start_time

        # Print summary
        self._print_summary(results, total_time)

        # Save and visualize results
        if SAVE_RESULTS:
            self._save_results(results)
            self._save_individual_images(results)  # Save each image with bounding boxes

        if SHOW_PLOTS:
            #self._visualize_results(results)
            self._create_comprehensive_summary(results)  # Create complete grid summary

        return results

    def _print_summary(self, results, total_time):
        """Print testing summary"""
        print(f"\n[SUMMARY] TESTING SUMMARY")
        print(f"="*40)

        total_images = len(results)
        successful = len([r for r in results if r['status'] == 'success'])
        total_faces = sum(r['num_faces'] for r in results)
        no_faces = len([r for r in results if r['status'] == 'no_faces'])
        errors = len([r for r in results if r['status'] == 'error'])

        print(f"Images processed: {total_images}")
        print(f"Successful: {successful}")
        print(f"No faces detected: {no_faces}")
        print(f"Errors: {errors}")
        print(f"Total faces detected: {total_faces}")
        print(f"Processing time: {total_time:.1f}s")
        print(f"Average time per image: {total_time/total_images:.2f}s")

        # Classification summary
        if total_faces > 0:
            real_faces = sum(len([f for f in r['faces'] if f['classification']['prediction'] == 'REAL'])
                           for r in results if r['status'] == 'success')
            fake_faces = total_faces - real_faces

            print(f"\n[RESULTS] Classification Results:")
            print(f"Real faces: {real_faces} ({real_faces/total_faces:.1%})")
            print(f"Fake faces: {fake_faces} ({fake_faces/total_faces:.1%})")

    def _save_results(self, results):
        """Save results to files"""
        os.makedirs(OUTPUT_PATH, exist_ok=True)

        # Save summary CSV
        import csv
        csv_path = os.path.join(OUTPUT_PATH, 'classification_results.csv')

        with open(csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['Image', 'Face_ID', 'Prediction', 'Confidence', 'Probability', 'Bbox_X', 'Bbox_Y', 'Bbox_W', 'Bbox_H'])

            for result in results:
                if result['status'] == 'success':
                    for face in result['faces']:
                        bbox = face['bbox']
                        cls = face['classification']
                        writer.writerow([
                            result['image_name'],
                            face['face_id'],
                            cls['prediction'],
                            f"{cls['confidence']:.3f}",
                            f"{cls['probability']:.3f}",
                            bbox[0], bbox[1], bbox[2], bbox[3]
                        ])

        print(f"[SAVED] Results saved to: {csv_path}")

    def _save_individual_images(self, results):
        """Save each processed image with bounding boxes and classifications"""
        os.makedirs(OUTPUT_PATH, exist_ok=True)
        individual_dir = os.path.join(OUTPUT_PATH, 'annotated_images')
        os.makedirs(individual_dir, exist_ok=True)

        saved_count = 0
        for result in results:
            if result['status'] in ['success', 'no_faces']:
                try:
                    # Load original image
                    if 'image' in result:
                        image = result['image'].copy()
                    else:
                        image = Image.open(result['image_path']).convert('RGB')

                    # Draw bounding boxes and labels
                    draw = ImageDraw.Draw(image)

                    # Try to use a larger font
                    try:
                        font = ImageFont.truetype("arial.ttf", 20)
                    except:
                        font = ImageFont.load_default()

                    if result['num_faces'] > 0:
                        for face in result['faces']:
                            bbox = face['bbox']
                            cls = face['classification']
                            
                            # Choose color based on prediction
                            color = 'green' if cls['prediction'] == 'REAL' else 'red'
                            
                            # Draw bounding box
                            x, y, w, h = bbox
                            draw.rectangle([x, y, x+w, y+h], outline=color, width=3)
                            
                            # Create label with prediction and confidence
                            label = f"{cls['prediction']} ({cls['confidence']:.1%})"
                            
                            # Draw label background
                            text_bbox = draw.textbbox((x, y-25), label, font=font)
                            draw.rectangle(text_bbox, fill=color)
                            
                            # Draw label text
                            draw.text((x, y-25), label, fill='white', font=font)
                    else:
                        # Add "NO FACES" label for images without faces
                        draw.text((10, 10), "NO FACES DETECTED", fill='orange', font=font)

                    # Save annotated image
                    base_name = os.path.splitext(result['image_name'])[0]
                    output_filename = f"{base_name}_annotated.jpg"
                    output_path = os.path.join(individual_dir, output_filename)
                    
                    image.save(output_path, 'JPEG', quality=95)
                    saved_count += 1

                except Exception as e:
                    print(f"[WARNING] Could not save annotated image for {result['image_name']}: {e}")

        print(f"[SAVED] {saved_count} annotated images saved to: {individual_dir}")

    def _visualize_results(self, results, max_display=6):
        """Visualize results with matplotlib (limited sample)"""
        try:
            # Filter successful results with faces
            display_results = [r for r in results if r['status'] == 'success' and r['num_faces'] > 0]
            display_results = display_results[:max_display]

            if not display_results:
                print("No results to visualize")
                return

            # Create subplots
            fig, axes = plt.subplots(2, 3, figsize=(15, 10))
            axes = axes.flatten()

            for i, result in enumerate(display_results):
                if i >= len(axes):
                    break

                ax = axes[i]
                image = result['image']

                # Draw bounding boxes on image
                draw_image = image.copy()
                draw = ImageDraw.Draw(draw_image)

                for face in result['faces']:
                    bbox = face['bbox']
                    cls = face['classification']

                    # Choose color based on prediction
                    color = 'green' if cls['prediction'] == 'REAL' else 'red'

                    # Draw bounding box
                    draw.rectangle([bbox[0], bbox[1], bbox[0]+bbox[2], bbox[1]+bbox[3]],
                                 outline=color, width=3)

                    # Add label
                    label = f"{cls['prediction']} ({cls['confidence']:.1%})"
                    draw.text((bbox[0], bbox[1]-20), label, fill=color)

                # Display image
                ax.imshow(draw_image)
                ax.set_title(f"{result['image_name']}\n{result['num_faces']} face(s)")
                ax.axis('off')

            # Hide empty subplots
            for i in range(len(display_results), len(axes)):
                axes[i].axis('off')

            plt.tight_layout()
            plt.savefig(os.path.join(OUTPUT_PATH, 'sample_classification.png'), dpi=150, bbox_inches='tight')
            plt.show()

        except Exception as e:
            print(f"[WARNING] Could not create sample visualization: {e}")

    def _create_comprehensive_summary(self, results):
        """Create a comprehensive grid summary of all processed images"""
        try:
            # Include all results (successful, no_faces, errors)
            all_results = results
            
            if not all_results:
                print("No results to create comprehensive summary")
                return

            # Calculate grid dimensions
            total_images = len(all_results)
            cols = 4  # 4 images per row
            rows = (total_images + cols - 1) // cols  # Ceiling division
            
            # Create large figure
            fig, axes = plt.subplots(rows, cols, figsize=(20, 5*rows))
            
            # Handle single row case
            if rows == 1:
                axes = axes.reshape(1, -1) if total_images > 1 else [axes]
            
            # Flatten axes for easier indexing
            axes_flat = axes.flatten() if total_images > 1 else [axes]

            for i, result in enumerate(all_results):
                ax = axes_flat[i]
                
                try:
                    # Load image
                    if 'image' in result and result['image'] is not None:
                        image = result['image'].copy()
                    else:
                        image = Image.open(result['image_path']).convert('RGB')
                    
                    # Create annotated copy
                    draw_image = image.copy()
                    draw = ImageDraw.Draw(draw_image)
                    
                    # Set up title based on status
                    title_parts = [result['image_name'][:20]]  # Truncate long names
                    
                    if result['status'] == 'success':
                        if result['num_faces'] > 0:
                            # Draw faces with bounding boxes
                            for face in result['faces']:
                                bbox = face['bbox']
                                cls = face['classification']
                                
                                # Choose color
                                color = 'green' if cls['prediction'] == 'REAL' else 'red'
                                
                                # Draw bounding box
                                x, y, w, h = bbox
                                draw.rectangle([x, y, x+w, y+h], outline=color, width=2)
                                
                                # Add small label
                                label = f"{cls['prediction']}\n{cls['confidence']:.0%}"
                                draw.text((x, y-15), label, fill=color)
                            
                            title_parts.append(f"{result['num_faces']} face(s)")
                            
                            # Count real vs fake
                            real_count = sum(1 for f in result['faces'] if f['classification']['prediction'] == 'REAL')
                            fake_count = result['num_faces'] - real_count
                            if real_count > 0:
                                title_parts.append(f"Real: {real_count}")
                            if fake_count > 0:
                                title_parts.append(f"Fake: {fake_count}")
                        else:
                            title_parts.append("No faces")
                            # Add text overlay
                            draw.text((10, 10), "NO FACES", fill='orange')
                    
                    elif result['status'] == 'no_faces':
                        title_parts.append("No faces detected")
                        draw.text((10, 10), "NO FACES", fill='orange')
                    
                    elif result['status'] == 'error':
                        title_parts.append("Error")
                        draw.text((10, 10), "ERROR", fill='red')
                    
                    # Display image
                    ax.imshow(draw_image)
                    ax.set_title('\n'.join(title_parts), fontsize=8)
                    ax.axis('off')
                    
                except Exception as e:
                    # Handle individual image errors
                    ax.text(0.5, 0.5, f"Error loading\n{result['image_name']}", 
                           ha='center', va='center', transform=ax.transAxes)
                    ax.set_title(f"Error: {result['image_name'][:20]}")
                    ax.axis('off')
            
            # Hide unused subplots
            for i in range(total_images, len(axes_flat)):
                axes_flat[i].axis('off')
            
            # Add overall title with summary stats
            total_faces = sum(r['num_faces'] for r in results if r['status'] == 'success')
            real_faces = sum(len([f for f in r['faces'] if f['classification']['prediction'] == 'REAL']) 
                           for r in results if r['status'] == 'success')
            fake_faces = total_faces - real_faces
            
            fig.suptitle(f"Face Classification Results - {total_images} Images, {total_faces} Faces\n"
                        f"Real: {real_faces} ({real_faces/total_faces*100 if total_faces > 0 else 0:.1f}%), "
                        f"Fake: {fake_faces} ({fake_faces/total_faces*100 if total_faces > 0 else 0:.1f}%)", 
                        fontsize=16, y=0.98)
            
            plt.tight_layout()
            plt.subplots_adjust(top=0.92)
            
            # Save comprehensive summary
            summary_path = os.path.join(OUTPUT_PATH, 'comprehensive_summary.png')
            plt.savefig(summary_path, dpi=200, bbox_inches='tight')
            print(f"[SAVED] Comprehensive summary saved to: {summary_path}")
            
            plt.show()

        except Exception as e:
            print(f"[WARNING] Could not create comprehensive summary: {e}")

# =============================================================================
# MAIN TESTING FUNCTION
# =============================================================================

def main():
    """Main testing function"""
    print("[*] FACE CLASSIFIER TESTING")
    print("="*50)

    # Check if model exists
    if not os.path.exists(MODEL_PATH):
        print(f"[ERROR] Model file not found: {MODEL_PATH}")
        print("Please make sure you have trained the model first.")
        print("Expected file: best_face_classifier_real_data.pth")
        return

    # Check if test folder exists
    if not os.path.exists(TEST_IMAGES_PATH):
        print(f"[ERROR] Test images folder not found: {TEST_IMAGES_PATH}")
        print("Please check the path and make sure it contains images.")
        return

    try:
        # Initialize tester
        tester = FaceClassifierTester(MODEL_PATH)

        # Run tests
        results = tester.test_folder(TEST_IMAGES_PATH, max_images=20)  # Limit for demo

        print(f"\n[OK] Testing completed!")
        print(f"Check '{OUTPUT_PATH}' folder for detailed results.")

    except Exception as e:
        print(f"[ERROR] Testing failed: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()