#!/usr/bin/env python3
"""
Debug Training Script
This script helps debug the training process to see what's happening with image loading.
"""

import cv2
import numpy as np
import os
from livestream import preprocess_face

def debug_user_images(user_id):
    """Debug the images for a specific user"""
    print(f"🔍 Debugging User {user_id} images...")
    
    user_folder = os.path.join("registered_faces", user_id)
    if not os.path.exists(user_folder):
        print(f"❌ User folder {user_folder} not found!")
        return
    
    # Check structure
    angle_dirs = ['front', 'left', 'right', 'up', 'down']
    has_proper_structure = all(os.path.isdir(os.path.join(user_folder, angle)) for angle in angle_dirs)
    
    print(f"📁 Structure: {'✅ Proper' if has_proper_structure else '❌ Flat'}")
    
    if has_proper_structure:
        total_images = 0
        for angle in angle_dirs:
            angle_path = os.path.join(user_folder, angle)
            if os.path.exists(angle_path):
                images = [f for f in os.listdir(angle_path) if f.lower().endswith('.jpg')]
                print(f"  {angle}: {len(images)} images")
                total_images += len(images)
                
                # Test first image loading
                if images:
                    first_image = os.path.join(angle_path, images[0])
                    print(f"    Testing {first_image}...")
                    
                    # Try to load the image
                    img = cv2.imread(first_image, cv2.IMREAD_GRAYSCALE)
                    if img is not None:
                        print(f"      ✅ Loaded: {img.shape}, min: {img.min()}, max: {img.max()}")
                        
                        # Test preprocessing
                        try:
                            processed = preprocess_face(img)
                            print(f"      ✅ Preprocessed: {processed.shape}, min: {processed.min()}, max: {processed.max()}")
                        except Exception as e:
                            print(f"      ❌ Preprocessing failed: {e}")
                    else:
                        print(f"      ❌ Failed to load image")
        
        print(f"📊 Total images: {total_images}")
    else:
        # Flat structure
        images = [f for f in os.listdir(user_folder) if f.lower().endswith('.jpg')]
        print(f"📸 Flat structure: {len(images)} images")
        
        if images:
            first_image = os.path.join(user_folder, images[0])
            print(f"  Testing {first_image}...")
            
            img = cv2.imread(first_image, cv2.IMREAD_GRAYSCALE)
            if img is not None:
                print(f"    ✅ Loaded: {img.shape}, min: {img.min()}, max: {img.max()}")
                
                try:
                    processed = preprocess_face(img)
                    print(f"      ✅ Preprocessed: {processed.shape}, min: {processed.min()}, max: {processed.max()}")
                except Exception as e:
                    print(f"      ❌ Preprocessing failed: {e}")
            else:
                print(f"    ❌ Failed to load image")

def test_training_function():
    """Test the training function step by step"""
    print("\n🧪 Testing Training Function...")
    
    faces = []
    labels = []
    
    # Test loading images from User 332 (Alejo)
    user_id = "332"
    user_folder = os.path.join("registered_faces", user_id)
    
    if not os.path.exists(user_folder):
        print(f"❌ User folder {user_folder} not found!")
        return
    
    angle_dirs = ['front', 'left', 'right', 'up', 'down']
    
    for angle in angle_dirs:
        angle_folder = os.path.join(user_folder, angle)
        if os.path.exists(angle_folder):
            print(f"📁 Processing {angle} angle...")
            
            for filename in os.listdir(angle_folder):
                if filename.endswith('.jpg'):
                    img_path = os.path.join(angle_folder, filename)
                    print(f"  📸 Loading {filename}...")
                    
                    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
                    if img is None:
                        print(f"    ❌ Failed to load {filename}")
                        continue
                    
                    print(f"    ✅ Loaded: {img.shape}")
                    
                    try:
                        img = preprocess_face(img)
                        faces.append(img)
                        labels.append(int(user_id))
                        print(f"    ✅ Preprocessed and added to training data")
                    except Exception as e:
                        print(f"    ❌ Preprocessing failed: {e}")
                        continue
    
    print(f"\n📊 Training Data Summary:")
    print(f"  Total faces: {len(faces)}")
    print(f"  Total labels: {len(labels)}")
    print(f"  Unique labels: {len(set(labels))}")
    
    if faces:
        print(f"  Face shape: {faces[0].shape}")
        print(f"  Face data type: {faces[0].dtype}")
        print(f"  Face value range: {faces[0].min()} to {faces[0].max()}")
    
    return faces, labels

def main():
    """Main debug function"""
    print("Debug Training Process")
    print("=" * 30)
    
    # Debug User 332 (Alejo)
    debug_user_images("332")
    
    # Test training function
    faces, labels = test_training_function()
    
    if faces and labels:
        print(f"\n✅ Training data looks good!")
        print(f"   Ready to train model with {len(faces)} images")
    else:
        print(f"\n❌ Training data has issues!")
        print(f"   Check the errors above")

if __name__ == "__main__":
    main()