#!/usr/bin/env python
"""
Atomizer Best Iteration Finder

Finds the iteration folder containing the best design for a study.
Useful for extracting the best model files after optimization.

Usage:
    python tools/find_best_iteration.py <study_name>
    python tools/find_best_iteration.py m1_mirror_adaptive_V14
    python tools/find_best_iteration.py m1_mirror_adaptive_V14 --copy-to ./best_design

Author: Atomizer
Created: 2025-12-12
"""

import argparse
import json
import shutil
import sqlite3
from pathlib import Path
import sys


def find_study_path(study_name: str) -> Path:
    """Find study directory by name."""
    studies_dir = Path(__file__).parent.parent / "studies"
    study_path = studies_dir / study_name

    if not study_path.exists():
        raise FileNotFoundError(f"Study not found: {study_path}")

    return study_path


def get_best_trial(study_path: Path) -> dict:
    """Get the best trial from the database."""
    db_path = study_path / "3_results" / "study.db"
    if not db_path.exists():
        raise FileNotFoundError(f"Database not found: {db_path}")

    conn = sqlite3.connect(str(db_path))
    cursor = conn.cursor()

    # Get trial with best weighted_sum (or lowest objective 0 value)
    cursor.execute("""
        SELECT t.number, tua_ws.value_json, tua_src.value_json
        FROM trials t
        LEFT JOIN trial_user_attributes tua_ws ON t.trial_id = tua_ws.trial_id AND tua_ws.key = 'weighted_sum'
        LEFT JOIN trial_user_attributes tua_src ON t.trial_id = tua_src.trial_id AND tua_src.key = 'source'
        WHERE t.state = 'COMPLETE'
        AND CAST(tua_ws.value_json AS REAL) < 1000
        ORDER BY CAST(tua_ws.value_json AS REAL) ASC
        LIMIT 1
    """)

    result = cursor.fetchone()
    if not result:
        # Fallback to trial_values
        cursor.execute("""
            SELECT t.number, tv.value
            FROM trials t
            JOIN trial_values tv ON t.trial_id = tv.trial_id AND tv.objective = 0
            WHERE t.state = 'COMPLETE'
            AND tv.value < 1000
            ORDER BY tv.value ASC
            LIMIT 1
        """)
        result = cursor.fetchone()
        if result:
            conn.close()
            return {
                "number": result[0],
                "weighted_sum": result[1],
                "source": "unknown"
            }
        raise ValueError("No valid trials found")

    conn.close()
    return {
        "number": result[0],
        "weighted_sum": float(result[1]) if result[1] else None,
        "source": json.loads(result[2]) if result[2] else "unknown"
    }


def trial_to_iteration(study_path: Path, trial_number: int) -> int:
    """Convert trial number to iteration number."""
    db_path = study_path / "3_results" / "study.db"
    conn = sqlite3.connect(str(db_path))
    cursor = conn.cursor()

    # Get all FEA trial numbers in order
    cursor.execute("""
        SELECT t.number
        FROM trials t
        JOIN trial_user_attributes tua ON t.trial_id = tua.trial_id
        WHERE t.state = 'COMPLETE' AND tua.key = 'source' AND tua.value_json = '"FEA"'
        ORDER BY t.number
    """)
    fea_trials = [r[0] for r in cursor.fetchall()]
    conn.close()

    if trial_number not in fea_trials:
        raise ValueError(f"Trial {trial_number} is not an FEA trial (seeded data has no iteration folder)")

    return fea_trials.index(trial_number) + 1


def find_best_iteration(study_name: str) -> dict:
    """Find the best iteration folder for a study."""
    study_path = find_study_path(study_name)
    best_trial = get_best_trial(study_path)

    result = {
        "study_name": study_name,
        "trial_number": best_trial["number"],
        "weighted_sum": best_trial["weighted_sum"],
        "source": best_trial["source"],
        "iteration_folder": None,
        "iteration_path": None
    }

    # Only FEA trials have iteration folders
    if best_trial["source"] == "FEA":
        iter_num = trial_to_iteration(study_path, best_trial["number"])
        iter_folder = f"iter{iter_num}"
        iter_path = study_path / "2_iterations" / iter_folder

        if iter_path.exists():
            result["iteration_folder"] = iter_folder
            result["iteration_path"] = str(iter_path)

    return result


def copy_best_iteration(study_name: str, dest_path: str):
    """Copy the best iteration folder to a destination."""
    result = find_best_iteration(study_name)

    if not result["iteration_path"]:
        raise ValueError(f"Best trial ({result['trial_number']}) is from seeded data, no iteration folder")

    src = Path(result["iteration_path"])
    dst = Path(dest_path)

    if dst.exists():
        print(f"Removing existing destination: {dst}")
        shutil.rmtree(dst)

    print(f"Copying {src} to {dst}")
    shutil.copytree(src, dst)

    return result


def main():
    parser = argparse.ArgumentParser(description="Find best iteration folder for Atomizer study")
    parser.add_argument("study_name", help="Name of the study")
    parser.add_argument("--copy-to", "-c", help="Copy best iteration folder to this path")
    parser.add_argument("--json", "-j", action="store_true", help="Output as JSON")

    args = parser.parse_args()

    try:
        if args.copy_to:
            result = copy_best_iteration(args.study_name, args.copy_to)
            print(f"\nCopied best iteration to: {args.copy_to}")
        else:
            result = find_best_iteration(args.study_name)

        if args.json:
            print(json.dumps(result, indent=2))
        else:
            print(f"\n{'='*60}")
            print(f"BEST ITERATION FOR: {result['study_name']}")
            print(f"{'='*60}")
            print(f"  Trial Number:      {result['trial_number']}")
            print(f"  Weighted Sum:      {result['weighted_sum']:.2f}" if result['weighted_sum'] else "  Weighted Sum:      N/A")
            print(f"  Source:            {result['source']}")
            print(f"  Iteration Folder:  {result['iteration_folder'] or 'N/A (seeded data)'}")
            if result['iteration_path']:
                print(f"  Full Path:         {result['iteration_path']}")
            print(f"{'='*60}")

    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()