Atomizer/tools/migrate_to_spec_v2.py

#!/usr/bin/env python
"""
AtomizerSpec v2.0 Migration CLI Tool

Migrates legacy optimization_config.json files to the new AtomizerSpec v2.0 format.

Usage:
    python tools/migrate_to_spec_v2.py studies/M1_Mirror/study_name
    python tools/migrate_to_spec_v2.py --all                       # Migrate all studies
    python tools/migrate_to_spec_v2.py --dry-run studies/*         # Preview without saving

Options:
    --dry-run       Preview migration without saving files
    --validate      Validate output against schema
    --all           Migrate all studies in studies/ directory
    --force         Overwrite existing atomizer_spec.json files
    --verbose       Show detailed migration info
"""

import argparse
import json
import sys
from pathlib import Path
from typing import List, Optional

# Add project root to path
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from optimization_engine.config.migrator import SpecMigrator, MigrationError
from optimization_engine.config.spec_validator import SpecValidator, SpecValidationError


def find_config_file(study_path: Path) -> Optional[Path]:
    """Find the optimization_config.json for a study."""
    # Check common locations
    candidates = [
        study_path / "1_setup" / "optimization_config.json",
        study_path / "optimization_config.json",
    ]

    for path in candidates:
        if path.exists():
            return path

    return None


def find_all_studies(studies_dir: Path) -> List[Path]:
    """Find all study directories with config files."""
    studies = []

    for item in studies_dir.rglob("optimization_config.json"):
        # Skip archives
        if "_archive" in str(item) or "archive" in str(item).lower():
            continue

        # Get study directory
        if item.parent.name == "1_setup":
            study_dir = item.parent.parent
        else:
            study_dir = item.parent

        if study_dir not in studies:
            studies.append(study_dir)

    return sorted(studies)


def migrate_study(
    study_path: Path,
    dry_run: bool = False,
    validate: bool = True,
    force: bool = False,
    verbose: bool = False
) -> bool:
    """
    Migrate a single study.

    Returns True if successful, False otherwise.
    """
    study_path = Path(study_path)

    if not study_path.exists():
        print(f"  ERROR: Study path does not exist: {study_path}")
        return False

    # Find config file
    config_path = find_config_file(study_path)
    if not config_path:
        print(f"  SKIP: No optimization_config.json found")
        return False

    # Check if spec already exists
    spec_path = study_path / "atomizer_spec.json"
    if spec_path.exists() and not force:
        print(f"  SKIP: atomizer_spec.json already exists (use --force to overwrite)")
        return False

    try:
        # Load old config
        with open(config_path, 'r', encoding='utf-8') as f:
            old_config = json.load(f)

        # Migrate
        migrator = SpecMigrator(study_path)
        new_spec = migrator.migrate(old_config)

        if verbose:
            print(f"  Config type: {migrator._detect_config_type(old_config)}")
            print(f"  Design variables: {len(new_spec['design_variables'])}")
            print(f"  Extractors: {len(new_spec['extractors'])}")
            print(f"  Objectives: {len(new_spec['objectives'])}")
            print(f"  Constraints: {len(new_spec.get('constraints', []))}")

        # Validate
        if validate:
            validator = SpecValidator()
            report = validator.validate(new_spec, strict=False)

            if not report.valid:
                print(f"  WARNING: Validation failed:")
                for err in report.errors[:3]:
                    print(f"    - {err.path}: {err.message}")
                if len(report.errors) > 3:
                    print(f"    ... and {len(report.errors) - 3} more errors")

        # Save
        if not dry_run:
            with open(spec_path, 'w', encoding='utf-8') as f:
                json.dump(new_spec, f, indent=2, ensure_ascii=False)
            print(f"  SUCCESS: Created {spec_path.name}")
        else:
            print(f"  DRY-RUN: Would create {spec_path.name}")

        return True

    except MigrationError as e:
        print(f"  ERROR: Migration failed: {e}")
        return False
    except json.JSONDecodeError as e:
        print(f"  ERROR: Invalid JSON in config: {e}")
        return False
    except Exception as e:
        print(f"  ERROR: Unexpected error: {e}")
        if verbose:
            import traceback
            traceback.print_exc()
        return False


def main():
    parser = argparse.ArgumentParser(
        description="Migrate optimization configs to AtomizerSpec v2.0",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__
    )

    parser.add_argument(
        "studies",
        nargs="*",
        help="Study directories to migrate (or use --all)"
    )
    parser.add_argument(
        "--all",
        action="store_true",
        help="Migrate all studies in studies/ directory"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Preview migration without saving files"
    )
    parser.add_argument(
        "--validate",
        action="store_true",
        default=True,
        help="Validate output against schema (default: True)"
    )
    parser.add_argument(
        "--no-validate",
        action="store_true",
        help="Skip validation"
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Overwrite existing atomizer_spec.json files"
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Show detailed migration info"
    )

    args = parser.parse_args()

    # Determine studies to migrate
    studies_dir = PROJECT_ROOT / "studies"

    if args.all:
        studies = find_all_studies(studies_dir)
        print(f"Found {len(studies)} studies to migrate\n")
    elif args.studies:
        studies = [Path(s) for s in args.studies]
    else:
        parser.print_help()
        return 1

    if not studies:
        print("No studies found to migrate")
        return 1

    # Migrate each study
    success_count = 0
    skip_count = 0
    error_count = 0

    for study_path in studies:
        # Handle relative paths
        if not study_path.is_absolute():
            # Try relative to CWD first, then project root
            if study_path.exists():
                pass
            elif (PROJECT_ROOT / study_path).exists():
                study_path = PROJECT_ROOT / study_path
            elif (studies_dir / study_path).exists():
                study_path = studies_dir / study_path

        print(f"Migrating: {study_path.name}")

        result = migrate_study(
            study_path,
            dry_run=args.dry_run,
            validate=not args.no_validate,
            force=args.force,
            verbose=args.verbose
        )

        if result:
            success_count += 1
        elif "SKIP" in str(result):
            skip_count += 1
        else:
            error_count += 1

    # Summary
    print(f"\n{'='*50}")
    print(f"Migration complete:")
    print(f"  Successful: {success_count}")
    print(f"  Skipped: {skip_count}")
    print(f"  Errors: {error_count}")

    if args.dry_run:
        print("\n(Dry run - no files were modified)")

    return 0 if error_count == 0 else 1


if __name__ == "__main__":
    sys.exit(main())