setup.py

#!/usr/bin/env python3
"""
Setup script for math-eval dataset generator.

This script sets up the environment and verifies everything is working correctly.
"""

import os
import sys
import subprocess
import platform

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

def run_command(cmd, description):
    """Run a command and handle errors."""
    print(f"🔄 {description}...")
    try:
        result = subprocess.run(cmd, shell=True, check=True, capture_output=True, text=True)
        print(f"✅ {description} completed successfully")
        return True
    except subprocess.CalledProcessError as e:
        print(f"❌ {description} failed:")
        print(f"   Error: {e}")
        if e.stderr:
            print(f"   Details: {e.stderr}")
        return False

def check_python_version():
    """Check if Python version is compatible."""
    print_header("Checking Python Version")
    
    version = sys.version_info
    print(f"Python version: {version.major}.{version.minor}.{version.micro}")
    
    if version.major < 3 or (version.major == 3 and version.minor < 7):
        print("❌ Python 3.7 or higher is required")
        return False
    
    print("✅ Python version is compatible")
    return True

def install_dependencies():
    """Install required dependencies."""
    print_header("Installing Dependencies")
    
    # Check if pip is available
    if not run_command("python -m pip --version", "Checking pip availability"):
        print("❌ pip is not available. Please install pip first.")
        return False
    
    # Install requirements
    if os.path.exists("requirements.txt"):
        return run_command("python -m pip install -r requirements.txt", "Installing requirements")
    else:
        print("❌ requirements.txt not found")
        return False

def check_directory_structure():
    """Check if required directories and files exist."""
    print_header("Checking Directory Structure")
    
    required_files = [
        "equation_generator.py",
        "verifier.py", 
        "generate_ocr_custom.py",
        "generate_visual_questions.py",
        "generate_counting_questions.py",
        "generate_partial_visual_questions.py",
        "run_pipeline.py",
        "requirements.txt"
    ]
    
    missing_files = []
    for file in required_files:
        if os.path.exists(file):
            print(f"✅ {file}")
        else:
            print(f"❌ {file} - Missing")
            missing_files.append(file)
    
    # Check for icon directory
    if os.path.exists("colored_icons_final"):
        print("✅ colored_icons_final/ directory found")
        
        # Count icons
        icon_count = 0
        for root, dirs, files in os.walk("colored_icons_final"):
            icon_count += len([f for f in files if f.endswith('.png')])
        
        print(f"   Found {icon_count} icon files")
        
        if icon_count == 0:
            print("⚠️  Warning: No icon files found in colored_icons_final/")
    else:
        print("⚠️  colored_icons_final/ directory not found")
        print("   Visual generation features will not work")
    
    # Check sample data
    sample_files = ["two-vars.txt", "three-vars.txt"]
    for file in sample_files:
        if os.path.exists(file):
            print(f"✅ {file}")
        else:
            print(f"⚠️  {file} - Not found (sample data)")
    
    return len(missing_files) == 0

def run_basic_tests():
    """Run basic functionality tests."""
    print_header("Running Basic Tests")
    
    # Test equation generation
    test_cmd = "python equation_generator.py --output_file test_setup.txt --num 3 --vars 2"
    if not run_command(test_cmd, "Testing equation generation"):
        return False
    
    # Test verification
    if os.path.exists("test_setup.txt"):
        verify_cmd = "python verifier.py --file test_setup.txt"
        if not run_command(verify_cmd, "Testing equation verification"):
            return False
        
        # Clean up
        os.remove("test_setup.txt")
    
    return True

def create_output_directories():
    """Create output directory structure."""
    print_header("Creating Output Directories")
    
    dirs = [
        "outputs",
        "outputs/equations",
        "outputs/visual",
        "outputs/visual/char_only",
        "outputs/visual/icon_only", 
        "outputs/visual/icon_partial",
        "outputs/counting",
        "outputs/logs"
    ]
    
    for dir_path in dirs:
        os.makedirs(dir_path, exist_ok=True)
        print(f"✅ Created {dir_path}")
    
    return True

def main():
    """Main setup function."""
    print_header("Math-Eval Setup")
    print("Welcome to the Math-Eval dataset generator setup!")
    print(f"Platform: {platform.system()} {platform.release()}")
    
    setup_steps = [
        ("Python Version Check", check_python_version),
        ("Dependency Installation", install_dependencies),
        ("Directory Structure Check", check_directory_structure),
        ("Output Directory Creation", create_output_directories),
        ("Basic Functionality Tests", run_basic_tests)
    ]
    
    failed_steps = []
    
    for step_name, step_func in setup_steps:
        if not step_func():
            failed_steps.append(step_name)
    
    print_header("Setup Summary")
    
    if not failed_steps:
        print("🎉 Setup completed successfully!")
        print("\nNext steps:")
        print("1. Generate your first dataset:")
        print("   python run_pipeline.py --num_equations 10 --task all")
        print("2. Read the README.md for detailed usage instructions")
        print("3. Run the test suite: python test_suite.py")
        
    else:
        print("❌ Setup completed with errors:")
        for step in failed_steps:
            print(f"   - {step}")
        print("\nPlease fix the issues above and run setup again.")
        sys.exit(1)

if __name__ == "__main__":
    main()