#!/usr/bin/env python
"""
Atomizer Best Design Archiver

Archives the best design iteration folder to 3_results/best_design_archive/
with timestamp. Should be called at the end of every optimization run.

Usage:
    python tools/archive_best_design.py <study_name>
    python tools/archive_best_design.py m1_mirror_adaptive_V14

    # Or call from run_optimization.py at the end:
    from tools.archive_best_design import archive_best_design
    archive_best_design("my_study")

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

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

# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))


def find_study_path(study_name: str) -> Path:
    """Find study directory by name."""
    # Check if already a path
    if Path(study_name).exists():
        return Path(study_name)

    # Check relative to tools dir
    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_info(study_path: Path) -> dict:
    """Get the best trial information 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
    cursor.execute("""
        SELECT t.number,
               tua_ws.value_json as weighted_sum,
               tua_40.value_json as obj_40,
               tua_60.value_json as obj_60,
               tua_mfg.value_json as obj_mfg,
               tua_src.value_json as source
        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_40 ON t.trial_id = tua_40.trial_id AND tua_40.key = 'rel_filtered_rms_40_vs_20'
        LEFT JOIN trial_user_attributes tua_60 ON t.trial_id = tua_60.trial_id AND tua_60.key = 'rel_filtered_rms_60_vs_20'
        LEFT JOIN trial_user_attributes tua_mfg ON t.trial_id = tua_mfg.trial_id AND tua_mfg.key = 'mfg_90_optician_workload'
        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 objective 0
        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 {
                "trial_number": result[0],
                "weighted_sum": result[1],
                "objectives": {},
                "source": "unknown"
            }
        raise ValueError("No valid trials found")

    conn.close()

    return {
        "trial_number": result[0],
        "weighted_sum": float(result[1]) if result[1] else None,
        "objectives": {
            "rel_filtered_rms_40_vs_20": float(result[2]) if result[2] else None,
            "rel_filtered_rms_60_vs_20": float(result[3]) if result[3] else None,
            "mfg_90_optician_workload": float(result[4]) if result[4] else None,
        },
        "source": json.loads(result[5]) if result[5] 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()

    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 archive_best_design(study_name: str, verbose: bool = True) -> dict:
    """
    Archive the best design to 3_results/best_design_archive/<timestamp>/

    Args:
        study_name: Name of the study or path to study directory
        verbose: Print progress messages

    Returns:
        Dictionary with archive info
    """
    study_path = find_study_path(study_name)

    # Get best trial info
    best_trial = get_best_trial_info(study_path)

    if best_trial["source"] != "FEA":
        if verbose:
            print(f"[WARN] Best trial #{best_trial['trial_number']} is from seeded data ({best_trial['source']})")
            print("[WARN] No iteration folder to archive. Skipping.")
        return {
            "success": False,
            "reason": "Best trial is from seeded data, no iteration folder",
            "trial_number": best_trial["trial_number"],
            "source": best_trial["source"]
        }

    # Find iteration folder
    iter_num = trial_to_iteration(study_path, best_trial["trial_number"])
    iter_folder = f"iter{iter_num}"
    iter_path = study_path / "2_iterations" / iter_folder

    if not iter_path.exists():
        raise FileNotFoundError(f"Iteration folder not found: {iter_path}")

    # Create archive directory with timestamp
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    archive_dir = study_path / "3_results" / "best_design_archive" / timestamp
    archive_dir.mkdir(parents=True, exist_ok=True)

    if verbose:
        print(f"[INFO] Archiving best design...")
        print(f"       Trial: #{best_trial['trial_number']}")
        print(f"       Iteration: {iter_folder}")
        print(f"       Weighted Sum: {best_trial['weighted_sum']:.2f}")
        print(f"       Archive: {archive_dir}")

    # Copy all files from iteration folder
    files_copied = 0
    for src_file in iter_path.iterdir():
        dst_file = archive_dir / src_file.name
        if src_file.is_file():
            shutil.copy2(src_file, dst_file)
            files_copied += 1
        elif src_file.is_dir():
            shutil.copytree(src_file, dst_file)
            files_copied += 1

    # Write metadata file
    metadata = {
        "study_name": study_path.name,
        "trial_number": best_trial["trial_number"],
        "iteration_folder": iter_folder,
        "weighted_sum": best_trial["weighted_sum"],
        "objectives": best_trial["objectives"],
        "source": best_trial["source"],
        "archived_at": datetime.now().isoformat(),
        "files_copied": files_copied
    }

    metadata_path = archive_dir / "_archive_info.json"
    with open(metadata_path, "w") as f:
        json.dump(metadata, f, indent=2)

    if verbose:
        print(f"[OK] Archived {files_copied} files to {archive_dir.name}/")

    return {
        "success": True,
        "archive_path": str(archive_dir),
        "trial_number": best_trial["trial_number"],
        "iteration_folder": iter_folder,
        "weighted_sum": best_trial["weighted_sum"],
        "files_copied": files_copied
    }


def main():
    parser = argparse.ArgumentParser(
        description="Archive best design iteration folder",
        epilog="This should be run at the end of every optimization."
    )
    parser.add_argument("study_name", help="Name of the study to archive")
    parser.add_argument("--quiet", "-q", action="store_true", help="Suppress output")
    parser.add_argument("--json", "-j", action="store_true", help="Output as JSON")

    args = parser.parse_args()

    try:
        result = archive_best_design(args.study_name, verbose=not args.quiet)

        if args.json:
            print(json.dumps(result, indent=2))
        elif result["success"]:
            print(f"\n[SUCCESS] Best design archived to: {result['archive_path']}")

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


if __name__ == "__main__":
    main()