setup_and_train.py

#!/usr/bin/env python
"""
VintageGAN: One-Click Setup and Training

Interactive script that handles everything after dataset download.
Guides user through the complete workflow with minimal input.

Usage:
    python setup_and_train.py

Author: VintageGAN Project
Date: 2024
"""

import subprocess
import sys
from pathlib import Path


def print_header(text: str):
    """Print a formatted header."""
    print(f"\n{'='*60}")
    print(f"  {text}")
    print(f"{'='*60}\n")


def print_step(step_num: int, total: int, description: str):
    """Print step information."""
    print(f"\n[Step {step_num}/{total}] {description}")
    print("-" * 60)


def ask_yes_no(question: str, default: bool = True) -> bool:
    """Ask a yes/no question."""
    default_str = "Y/n" if default else "y/N"
    response = input(f"{question} [{default_str}]: ").strip().lower()

    if not response:
        return default
    return response in ["y", "yes"]


def check_dataset() -> tuple:
    """Check dataset status."""
    train_dir = Path("data/imagenet_subset/train")
    val_dir = Path("data/imagenet_subset/val")

    train_exists = train_dir.exists()
    val_exists = val_dir.exists()

    if train_exists and val_exists:
        train_images = list(train_dir.glob("**/*.jpg")) + list(
            train_dir.glob("**/*.png")
        )
        val_images = list(val_dir.glob("**/*.jpg")) + list(val_dir.glob("**/*.png"))
        return True, len(train_images), len(val_images)

    return False, 0, 0


def main():
    print_header("VintageGAN: Automated Setup and Training")

    print("This script will guide you through:")
    print("  1. Dataset verification")
    print("  2. Installation verification")
    print("  3. Generator pretraining (~6 hours)")
    print("  4. GAN training (~4 hours)")
    print("  5. Evaluation and results generation")
    print(
        "\nTotal time depends on dataset size, profile, and GPU. Start with local_256."
    )

    if not ask_yes_no("\nReady to begin?"):
        print("\nExiting. Run this script again when ready.")
        return 0

    # Step 1: Check dataset
    print_step(1, 6, "Dataset Verification")

    dataset_ready, train_count, val_count = check_dataset()

    if dataset_ready:
        print(f"✅ Dataset found:")
        print(f"   Training images: {train_count}")
        print(f"   Validation images: {val_count}")
    else:
        print("❌ Dataset not found")
        print("\nPlease download the dataset first:")
        print("  Option 1: Real data (recommended)")
        print(
            "    python training/download_data.py --mode full --num-train 10000 --num-val 1000"
        )
        print("\n  Option 2: Dummy data (for testing)")
        print("    python training/download_data.py --mode dummy")

        if ask_yes_no("\nCreate dummy dataset now for testing?"):
            print("\nCreating dummy dataset...")
            subprocess.run(
                [
                    sys.executable,
                    "-c",
                    "from training.download_data import create_dummy_dataset; "
                    "create_dummy_dataset('data/imagenet_subset', 100, 20)",
                ]
            )
            print("✅ Dummy dataset created")
        else:
            print("\nPlease download dataset and run this script again.")
            return 1

    # Step 2: Verify installation
    print_step(2, 6, "Installation Verification")

    print("Checking installation...")
    result = subprocess.run([sys.executable, "verify_installation.py"])

    if result.returncode != 0:
        print("\n❌ Installation verification failed")
        print("Please fix issues and run this script again.")
        return 1

    print("✅ Installation verified")

    # Step 3: Check for existing checkpoints
    print_step(3, 6, "Checkpoint Status")

    pretrain_checkpoint = Path("checkpoints/generator_pretrain_best.pth")
    final_checkpoint = Path("checkpoints/generator_final.pth")

    skip_pretrain = False
    skip_gan = False

    if pretrain_checkpoint.exists():
        print(f"✅ Pretrain checkpoint found: {pretrain_checkpoint}")
        if ask_yes_no("Skip pretraining and use existing checkpoint?"):
            skip_pretrain = True
    else:
        print("❌ No pretrain checkpoint found")

    if final_checkpoint.exists():
        print(f"✅ Final checkpoint found: {final_checkpoint}")
        if ask_yes_no("Skip GAN training and use existing checkpoint?"):
            skip_gan = True
    else:
        print("❌ No final checkpoint found")

    # Step 4: Estimate time
    print_step(4, 6, "Time Estimation")

    estimated_hours = 0
    if not skip_pretrain:
        estimated_hours += 6
        print("  - Generator pretraining: ~6 hours")
    if not skip_gan:
        estimated_hours += 4
        print("  - GAN training: ~4 hours")

    estimated_hours += 1  # Evaluation
    print("  - Evaluation: ~1 hour")

    print(f"\n  Total estimated time: ~{estimated_hours} hours")

    if estimated_hours > 1:
        print("\n⚠️  This will take a while. The script can run unattended.")
        print("   Progress will be logged to logs/ directory.")
        print("   You can monitor with: tail -f logs/*.log")

    if not ask_yes_no(f"\nProceed with ~{estimated_hours} hour training?"):
        print("\nTraining cancelled.")
        return 0

    # Step 5: Run automated pipeline
    print_step(5, 6, "Running Automated Training Pipeline")

    print("\n🚀 Starting automated pipeline...")
    print("   This will run unattended. You can close this window.")
    print("   Check logs/ directory for progress.")

    cmd = [sys.executable, "run_full_pipeline.py"]

    if skip_pretrain:
        cmd.append("--skip-pretrain")
    if skip_gan:
        cmd.append("--skip-gan")

    print(f"\nCommand: {' '.join(cmd)}\n")

    result = subprocess.run(cmd)

    if result.returncode != 0:
        print("\n❌ Pipeline failed")
        print("Check logs/ directory for details.")
        return 1

    # Step 6: Summary
    print_step(6, 6, "Training Complete!")

    print("✅ All training completed successfully!")
    print("\n📁 Generated files:")
    print("   Checkpoints: checkpoints/")
    print("   Results: results/")
    print("   Logs: logs/")

    print("\n📊 Results for your paper:")
    print("   - results/quantitative_metrics.json")
    print("   - results/results_table.csv (Excel/Sheets)")
    print("   - results/results_table.tex (LaTeX)")
    print("   - results/sample_results.png (Figure)")
    print("   - results/training_curves.png (Figure)")

    print("\n🎉 VintageGAN training complete!")
    print("   Use the generated results in your research paper.")

    # Offer to run inference demo
    if ask_yes_no("\nWould you like to test inference on a sample image?"):
        print("\nLaunching interactive demo...")
        subprocess.run(
            [sys.executable, "-m", "jupyter", "notebook", "notebooks/demo.ipynb"]
        )

    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        print("\n\n⚠️  Interrupted by user")
        print("You can resume by running this script again.")
        sys.exit(1)