#!/usr/bin/env python
"""
Migration script to reorganize studies into topic-based subfolders.

Run with --dry-run first to preview changes:
    python migrate_studies_to_topics.py --dry-run

Then run without flag to execute:
    python migrate_studies_to_topics.py
"""

import shutil
import argparse
from pathlib import Path

STUDIES_DIR = Path(__file__).parent / "studies"

# Topic classification based on study name prefixes
TOPIC_MAPPING = {
    'bracket_': 'Simple_Bracket',
    'drone_gimbal_': 'Drone_Gimbal',
    'm1_mirror_': 'M1_Mirror',
    'uav_arm_': 'UAV_Arm',
    'simple_beam_': 'Simple_Beam',
}

# Files/folders to skip (not studies)
SKIP_ITEMS = {
    'm1_mirror_all_trials_export.csv',  # Data export file
    '.gitkeep',
    '__pycache__',
}


def classify_study(study_name: str) -> str:
    """Determine which topic folder a study belongs to."""
    for prefix, topic in TOPIC_MAPPING.items():
        if study_name.startswith(prefix):
            return topic
    return '_Other'


def get_studies_to_migrate():
    """Get list of studies that need migration (not already in topic folders)."""
    studies = []

    for item in STUDIES_DIR.iterdir():
        # Skip non-directories and special items
        if not item.is_dir():
            continue
        if item.name in SKIP_ITEMS:
            continue
        if item.name.startswith('.'):
            continue

        # Check if this is already a topic folder (contains study subdirs)
        # A topic folder would have subdirs with 1_setup folders
        is_topic_folder = any(
            (sub / "1_setup").exists()
            for sub in item.iterdir()
            if sub.is_dir()
        )

        if is_topic_folder:
            print(f"[SKIP] {item.name} - already a topic folder")
            continue

        # Check if this is a study (has 1_setup or optimization_config.json)
        is_study = (
            (item / "1_setup").exists() or
            (item / "optimization_config.json").exists()
        )

        if is_study:
            topic = classify_study(item.name)
            studies.append({
                'name': item.name,
                'source': item,
                'topic': topic,
                'target': STUDIES_DIR / topic / item.name
            })
        else:
            print(f"[SKIP] {item.name} - not a study (no 1_setup folder)")

    return studies


def migrate_studies(dry_run: bool = True):
    """Migrate studies to topic folders."""
    studies = get_studies_to_migrate()

    if not studies:
        print("\nNo studies to migrate. All studies are already organized.")
        return

    # Group by topic for display
    by_topic = {}
    for s in studies:
        if s['topic'] not in by_topic:
            by_topic[s['topic']] = []
        by_topic[s['topic']].append(s)

    print("\n" + "="*60)
    print("MIGRATION PLAN")
    print("="*60)

    for topic in sorted(by_topic.keys()):
        print(f"\n{topic}/")
        for s in by_topic[topic]:
            print(f"  +-- {s['name']}/")

    print(f"\nTotal: {len(studies)} studies to migrate")

    if dry_run:
        print("\n[DRY RUN] No changes made. Run without --dry-run to execute.")
        return

    # Execute migration
    print("\n" + "="*60)
    print("EXECUTING MIGRATION")
    print("="*60)

    # Create topic folders
    created_topics = set()
    for s in studies:
        topic_dir = STUDIES_DIR / s['topic']
        if s['topic'] not in created_topics:
            topic_dir.mkdir(exist_ok=True)
            created_topics.add(s['topic'])
            print(f"[CREATE] {s['topic']}/")

    # Move studies
    for s in studies:
        try:
            shutil.move(str(s['source']), str(s['target']))
            print(f"[MOVE] {s['name']} -> {s['topic']}/{s['name']}")
        except Exception as e:
            print(f"[ERROR] Failed to move {s['name']}: {e}")

    print("\n" + "="*60)
    print("MIGRATION COMPLETE")
    print("="*60)


def main():
    parser = argparse.ArgumentParser(description="Migrate studies to topic folders")
    parser.add_argument('--dry-run', action='store_true',
                        help='Preview changes without executing')
    args = parser.parse_args()

    migrate_studies(dry_run=args.dry_run)


if __name__ == "__main__":
    main()