#!/usr/bin/env python3
"""
Quick Dataset Download Script for MorphGuard
Downloads a balanced real face dataset to improve training data ratio
Uses your actual API keys from unsplash_keys.json
"""
import os
import sys
sys.path.append('/root/MorphGuard')

from comprehensive_dataset_downloader import DatasetDownloader

def quick_balance_dataset():
    """Download datasets to balance the existing morph-heavy training data"""
    
    print("🚀 MorphGuard Quick Dataset Balancer")
    print("=" * 50)
    
    # Check current dataset balance
    print("📊 Current Dataset Analysis:")
    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
    real_count = len([f for f in os.listdir('data/train/real') if f.endswith('.jpg')]) if os.path.exists('data/train/real') else 0
    
    print(f"   Morph images: {morph_count:,}")
    print(f"   Real images: {real_count:,}")
    print(f"   Current ratio: {morph_count/max(real_count,1):.1f}:1 (morph:real)")
    
    # Target: Get to 3:1 ratio (need ~23,000 real images for 70,000 morphs)
    target_real = morph_count // 3
    needed = max(0, target_real - real_count)
    
    print(f"\n🎯 Target for balanced training:")
    print(f"   Target real images: {target_real:,}")
    print(f"   Need to download: {needed:,}")
    
    if needed <= 0:
        print("✅ Dataset already balanced!")
        return
    
    # Initialize downloader
    downloader = DatasetDownloader()
    
    # Download plan
    unsplash_count = min(3000, needed // 3)
    pexels_count = min(3000, needed // 3)
    lfw_count = 13000  # LFW has ~13k images
    generated_count = min(2000, needed // 4)
    
    print(f"\n📥 Download Plan:")
    print(f"   Unsplash: {unsplash_count:,} images")
    print(f"   Pexels: {pexels_count:,} images")
    print(f"   LFW Dataset: ~13,000 images")
    print(f"   AI-Generated: {generated_count:,} images")
    print(f"   Total planned: {unsplash_count + pexels_count + 13000 + generated_count:,}")
    
    input("\nPress Enter to start downloading or Ctrl+C to cancel...")
    
    total_downloaded = 0
    source_dirs = []
    
    # 1. Download LFW (academic standard)
    print("\n🔵 Phase 1: LFW Academic Dataset")
    lfw_dir = "data/real_faces/lfw"
    lfw_downloaded = downloader.download_lfw_dataset(lfw_dir)
    total_downloaded += lfw_downloaded
    source_dirs.append(os.path.join(lfw_dir, "lfw"))
    
    # 2. Download Unsplash (diverse, high-quality)
    print("\n🎨 Phase 2: Unsplash Professional Photos")
    unsplash_dir = "data/real_faces/unsplash"
    unsplash_downloaded = downloader.download_unsplash_faces(unsplash_count, unsplash_dir)
    total_downloaded += unsplash_downloaded
    source_dirs.append(unsplash_dir)
    
    # 3. Download Pexels (additional diversity)
    print("\n📷 Phase 3: Pexels Stock Photos")
    pexels_dir = "data/real_faces/pexels"
    pexels_downloaded = downloader.download_pexels_faces(pexels_count, pexels_dir)
    total_downloaded += pexels_downloaded
    source_dirs.append(pexels_dir)
    
    # 4. Download AI-generated (if still needed)
    remaining_needed = needed - total_downloaded
    if remaining_needed > 0 and generated_count > 0:
        print(f"\n🤖 Phase 4: AI-Generated Faces ({min(generated_count, remaining_needed):,} needed)")
        generated_dir = "data/real_faces/generated"
        generated_downloaded = downloader.download_generated_faces(min(generated_count, remaining_needed), generated_dir)
        total_downloaded += generated_downloaded
        source_dirs.append(generated_dir)
    
    # 5. Organize into training structure
    print("\n📁 Phase 5: Organizing for Training")
    train_count, val_count = downloader.organize_for_training(source_dirs)
    
    # Final summary
    print("\n" + "="*50)
    print("🎉 DOWNLOAD COMPLETE!")
    print("="*50)
    
    new_real_count = real_count + train_count
    new_ratio = morph_count / max(new_real_count, 1)
    
    print(f"📊 Updated Dataset Balance:")
    print(f"   Morph images: {morph_count:,}")
    print(f"   Real images: {new_real_count:,} (+{train_count:,})")
    print(f"   New ratio: {new_ratio:.1f}:1 (morph:real)")
    print(f"   Validation: +{val_count:,} real images")
    
    if new_ratio <= 4:
        print("✅ EXCELLENT: Dataset is now well-balanced for training!")
    elif new_ratio <= 6:
        print("🟡 GOOD: Dataset balance significantly improved!")
    else:
        print("🟠 BETTER: Some improvement, consider downloading more real images")
    
    print(f"\n📈 Ready to retrain with {total_downloaded:,} new real images!")

if __name__ == "__main__":
    # Change to MorphGuard directory
    os.chdir('/root/MorphGuard')
    quick_balance_dataset()