# preprocess.py
import cv2
import os
import numpy as np
from tqdm import tqdm
from mtcnn import MTCNN

print("Starting preprocessing with face detection...")

detector = MTCNN()

def process_images(input_dir, output_dir, label_name):
    os.makedirs(output_dir, exist_ok=True)
    
    files = os.listdir(input_dir)
    for file in tqdm(files, desc=f"Processing {label_name}"):
        if not file.endswith(('.jpg', '.png')):
            continue
            
        img_path = os.path.join(input_dir, file)
        img = cv2.imread(img_path)
        
        if img is None:
            continue
        
        # Detect faces
        rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        faces = detector.detect_faces(rgb)
        
        if len(faces) > 0:
            # Get largest face
            largest = max(faces, key=lambda x: x['box'][2] * x['box'][3])
            x, y, w, h = largest['box']
            
            # Add padding
            x = max(0, x - 20)
            y = max(0, y - 20)
            w = min(img.shape[1] - x, w + 40)
            h = min(img.shape[0] - y, h + 40)
            
            # Crop face
            face = img[y:y+h, x:x+w]
            face_resized = cv2.resize(face, (224, 224))
        else:
            # No face - use center crop
            h, w = img.shape[:2]
            face_resized = cv2.resize(img[h//4:3*h//4, w//4:3*w//4], (224, 224))
        
        cv2.imwrite(os.path.join(output_dir, file), face_resized)

# Process real and fake images
process_images("data/real", "processed/real", "REAL")
process_images("data/fake", "processed/fake", "FAKE")

print("✅ Preprocessing complete!")