Atomizer/reorganize_study.py

"""
Reorganize simple_beam_optimization study to new structure.
Handles locked files gracefully.
"""
import shutil
from pathlib import Path
import time

study_dir = Path("studies/simple_beam_optimization")

# Check current state
print("Current directory structure:")
print(f"  1_setup exists: {(study_dir / '1_setup').exists()}")
print(f"  2_substudies exists: {(study_dir / '2_substudies').exists()}")
print(f"  3_reports exists: {(study_dir / '3_reports').exists()}")
print()

# Copy full_optimization_50trials if not already done
src = study_dir / "substudies" / "full_optimization_50trials"
dst = study_dir / "2_substudies" / "04_full_optimization_50trials"

if src.exists() and not dst.exists():
    print(f"Copying {src.name} to {dst.name}...")
    try:
        shutil.copytree(src, dst)
        print(f"  SUCCESS: Copied to {dst}")
    except Exception as e:
        print(f"  WARNING: {e}")
        print(f"  Will attempt to continue...")

# Move OPTIMIZATION_RESULTS_50TRIALS.md
old_results_file = study_dir / "OPTIMIZATION_RESULTS_50TRIALS.md"
new_results_file = dst / "OPTIMIZATION_RESULTS.md"

if old_results_file.exists() and not new_results_file.exists():
    print(f"\nMoving {old_results_file.name}...")
    try:
        shutil.move(str(old_results_file), str(new_results_file))
        print(f"  SUCCESS: Moved to {new_results_file}")
    except Exception as e:
        print(f"  WARNING: {e}")

# Move COMPREHENSIVE_BENCHMARK_RESULTS.md
old_bench_file = study_dir / "COMPREHENSIVE_BENCHMARK_RESULTS.md"
new_bench_file = study_dir / "3_reports" / "COMPREHENSIVE_BENCHMARK_RESULTS.md"

if old_bench_file.exists() and not new_bench_file.exists():
    print(f"\nMoving {old_bench_file.name}...")
    try:
        shutil.move(str(old_bench_file), str(new_bench_file))
        print(f"  SUCCESS: Moved to {new_bench_file}")
    except Exception as e:
        print(f"  WARNING: {e}")

# Try to remove old substudies directory (may fail due to locked files - that's OK)
old_substudies = study_dir / "substudies"
if old_substudies.exists():
    print(f"\nAttempting to remove old 'substudies' directory...")
    try:
        # Try multiple times in case files get unlocked
        for attempt in range(3):
            try:
                shutil.rmtree(old_substudies)
                print(f"  SUCCESS: Removed old 'substudies' directory")
                break
            except Exception as e:
                if attempt < 2:
                    print(f"  Attempt {attempt + 1} failed, retrying in 1 second...")
                    time.sleep(1)
                else:
                    print(f"  INFO: Could not remove old 'substudies' directory (files may be locked)")
                    print(f"  You can manually delete it later: {old_substudies}")
    except Exception as e:
        print(f"  WARNING: {e}")

# Remove old model directory if empty
old_model = study_dir / "model"
if old_model.exists() and not list(old_model.iterdir()):
    print(f"\nRemoving empty 'model' directory...")
    try:
        old_model.rmdir()
        print(f"  SUCCESS: Removed empty 'model' directory")
    except Exception as e:
        print(f"  WARNING: {e}")

print("\n" + "="*70)
print("Reorganization complete!")
print("="*70)
print("\nNew structure:")
print("  1_setup/ - Pre-optimization setup")
print("  2_substudies/ - Optimization runs (numbered)")
print("  3_reports/ - Study-level analysis")
print()
print("Next steps:")
print("  1. Update study_metadata.json")
print("  2. Create substudy README files")
print("  3. Delete old 'substudies' folder manually if it still exists")