#!/usr/bin/env python3
"""
Integrate CelebA dataset with MorphGuard training
"""
import os
import shutil
import random
from pathlib import Path
from tqdm import tqdm

def integrate_celeba():
    """Integrate CelebA images into MorphGuard training pipeline"""
    
    print("🎭 Integrating CelebA with MorphGuard")
    print("=" * 50)
    
    # Check if CelebA exists
    celeba_dir = "data/raw/celeba/img_align_celeba"
    if not os.path.exists(celeba_dir):
        print(f"❌ CelebA directory not found: {celeba_dir}")
        print("Please copy your CelebA files first!")
        return
    
    # Count CelebA images
    celeba_images = [f for f in os.listdir(celeba_dir) if f.endswith('.jpg')]
    print(f"📊 Found {len(celeba_images):,} CelebA images")
    
    # Create output directories
    Path("data/celeba_processed/train").mkdir(parents=True, exist_ok=True)
    Path("data/celeba_processed/val").mkdir(parents=True, exist_ok=True)
    
    # Read evaluation partition (if available)
    partition_file = "data/raw/celeba/partitions/list_eval_partition.txt"
    if os.path.exists(partition_file):
        print("📋 Using official CelebA train/val/test split")
        with open(partition_file, 'r') as f:
            lines = f.readlines()
        
        train_images = []
        val_images = []
        
        for line in lines:
            if line.strip():
                img_name, partition = line.strip().split()
                if partition == '0':  # Training
                    train_images.append(img_name)
                elif partition == '1':  # Validation
                    val_images.append(img_name)
                # Skip test images (partition == '2')
    else:
        print("📋 Creating random train/val split (90/10)")
        random.shuffle(celeba_images)
        split_point = int(len(celeba_images) * 0.9)
        train_images = celeba_images[:split_point]
        val_images = celeba_images[split_point:]
    
    print(f"   Training: {len(train_images):,} images")
    print(f"   Validation: {len(val_images):,} images")
    
    # Copy training images (limit to prevent overload)
    max_train = 50000  # Reasonable limit
    max_val = 5000
    
    print("\n📂 Copying training images...")
    train_copied = 0
    for img_name in tqdm(train_images[:max_train], desc="Training"):
        src_path = os.path.join(celeba_dir, img_name)
        dst_path = os.path.join("data/celeba_processed/train", f"celeba_{img_name}")
        
        try:
            if os.path.exists(src_path):
                shutil.copy2(src_path, dst_path)
                train_copied += 1
        except Exception as e:
            continue
    
    print(f"📂 Copying validation images...")
    val_copied = 0
    for img_name in tqdm(val_images[:max_val], desc="Validation"):
        src_path = os.path.join(celeba_dir, img_name)
        dst_path = os.path.join("data/celeba_processed/val", f"celeba_{img_name}")
        
        try:
            if os.path.exists(src_path):
                shutil.copy2(src_path, dst_path)
                val_copied += 1
        except Exception as e:
            continue
    
    print(f"\n✅ CelebA Integration Complete:")
    print(f"   Training images: {train_copied:,}")
    print(f"   Validation images: {val_copied:,}")
    print(f"   Location: data/celeba_processed/")
    
    # Add to main training pipeline
    add_to_training_pipeline(train_copied, val_copied)

def add_to_training_pipeline(train_count, val_count):
    """Add CelebA images to main training pipeline"""
    print(f"\n🔄 Adding to main training pipeline...")
    
    # Copy to main training directories
    celeba_train_added = 0
    celeba_val_added = 0
    
    # Add to training set
    train_src = "data/celeba_processed/train"
    train_dst = "data/train/real"
    
    if os.path.exists(train_src) and os.path.exists(train_dst):
        for img_file in os.listdir(train_src):
            if img_file.endswith('.jpg'):
                src_path = os.path.join(train_src, img_file)
                dst_path = os.path.join(train_dst, img_file)
                
                if not os.path.exists(dst_path):
                    try:
                        shutil.copy2(src_path, dst_path)
                        celeba_train_added += 1
                    except Exception as e:
                        continue
    
    # Add to validation set
    val_src = "data/celeba_processed/val"
    val_dst = "data/val/real"
    
    if os.path.exists(val_src) and os.path.exists(val_dst):
        for img_file in os.listdir(val_src):
            if img_file.endswith('.jpg'):
                src_path = os.path.join(val_src, img_file)
                dst_path = os.path.join(val_dst, img_file)
                
                if not os.path.exists(dst_path):
                    try:
                        shutil.copy2(src_path, dst_path)
                        celeba_val_added += 1
                    except Exception as e:
                        continue
    
    # Calculate new dataset balance
    morph_count = len([f for f in os.listdir('data/train/morph') if f.endswith('.jpg')]) if os.path.exists('data/train/morph') else 0
    total_real_train = len([f for f in os.listdir('data/train/real') if f.endswith('.jpg')]) if os.path.exists('data/train/real') else 0
    total_real_val = len([f for f in os.listdir('data/val/real') if f.endswith('.jpg')]) if os.path.exists('data/val/real') else 0
    
    new_ratio = morph_count / max(total_real_train, 1)
    
    print(f"✅ Added to main pipeline:")
    print(f"   CelebA train added: {celeba_train_added:,}")
    print(f"   CelebA val added: {celeba_val_added:,}")
    print(f"\n📊 Updated Dataset Balance:")
    print(f"   Total morph: {morph_count:,}")
    print(f"   Total real train: {total_real_train:,}")
    print(f"   Total real val: {total_real_val:,}")
    print(f"   New ratio: {new_ratio:.1f}:1 (morph:real)")
    
    if new_ratio <= 3:
        print("🎯 EXCELLENT! Perfect balance for training!")
    elif new_ratio <= 5:
        print("🟢 VERY GOOD! Great balance for training!")
    else:
        print("🟡 IMPROVED! Better balance achieved!")

if __name__ == "__main__":
    integrate_celeba()