Atomizer/optimization_engine/utils/trial_manager.py

"""
Trial Manager - Unified trial numbering and folder management
==============================================================

Provides consistent trial_NNNN naming across all optimization methods
(Optuna, Turbo, GNN, manual) with proper database integration.

Usage:
    from optimization_engine.utils.trial_manager import TrialManager

    tm = TrialManager(study_dir)

    # Get next trial (creates folder, reserves DB row)
    trial = tm.new_trial(params={'rib_thickness': 10.5, ...})

    # After FEA completes
    tm.complete_trial(
        trial_id=trial['trial_id'],
        objectives={'wfe_40_20': 5.63, 'mass_kg': 118.67},
        metadata={'solve_time': 211.7}
    )

Key principles:
- Trial numbers NEVER reset (monotonically increasing)
- Folders NEVER get overwritten
- Database is always in sync with filesystem
- Surrogate predictions are NOT trials (only FEA results)
"""

import json
import sqlite3
import shutil
from pathlib import Path
from datetime import datetime
from typing import Dict, Any, Optional, List, Union
from filelock import FileLock

from .dashboard_db import DashboardDB


class TrialManager:
    """Manages trial numbering, folders, and database for optimization studies."""

    def __init__(self, study_dir: Union[str, Path], study_name: Optional[str] = None):
        """
        Initialize trial manager for a study.

        Args:
            study_dir: Path to study directory (contains 1_setup/, 2_iterations/, 3_results/)
            study_name: Name of study (defaults to directory name)
        """
        self.study_dir = Path(study_dir)
        self.study_name = study_name or self.study_dir.name

        self.iterations_dir = self.study_dir / "2_iterations"
        self.results_dir = self.study_dir / "3_results"
        self.db_path = self.results_dir / "study.db"
        self.lock_path = self.results_dir / ".trial_lock"

        # Ensure directories exist
        self.iterations_dir.mkdir(parents=True, exist_ok=True)
        self.results_dir.mkdir(parents=True, exist_ok=True)

        # Initialize database
        self.db = DashboardDB(self.db_path, self.study_name)

    def _get_next_trial_number(self) -> int:
        """Get next available trial number (never resets)."""
        # Check filesystem
        existing_folders = list(self.iterations_dir.glob("trial_*"))
        max_folder = 0
        for folder in existing_folders:
            try:
                num = int(folder.name.split('_')[1])
                max_folder = max(max_folder, num)
            except (IndexError, ValueError):
                continue

        # Check database
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute("SELECT COALESCE(MAX(number), -1) + 1 FROM trials")
        max_db = cursor.fetchone()[0]
        conn.close()

        # Return max of both + 1 (use 1-based for folders, 0-based for DB)
        return max(max_folder, max_db) + 1

    def new_trial(
        self,
        params: Dict[str, float],
        source: str = "turbo",
        metadata: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Start a new trial - creates folder and reserves DB row.

        Args:
            params: Design parameters for this trial
            source: How this trial was generated ("turbo", "optuna", "manual")
            metadata: Additional info (turbo_batch, predicted_ws, etc.)

        Returns:
            Dict with trial_id, trial_number, folder_path
        """
        # Use file lock to prevent race conditions
        with FileLock(self.lock_path):
            trial_number = self._get_next_trial_number()

            # Create folder with zero-padded name
            folder_name = f"trial_{trial_number:04d}"
            folder_path = self.iterations_dir / folder_name
            folder_path.mkdir(exist_ok=True)

            # Save params to folder
            params_file = folder_path / "params.json"
            with open(params_file, 'w') as f:
                json.dump(params, f, indent=2)

            # Also save as .exp format for NX compatibility
            exp_file = folder_path / "params.exp"
            with open(exp_file, 'w') as f:
                for name, value in params.items():
                    f.write(f"[mm]{name}={value}\n")

            # Save metadata
            meta = {
                "trial_number": trial_number,
                "source": source,
                "status": "RUNNING",
                "datetime_start": datetime.now().isoformat(),
                "params": params,
            }
            if metadata:
                meta.update(metadata)

            meta_file = folder_path / "_meta.json"
            with open(meta_file, 'w') as f:
                json.dump(meta, f, indent=2)

            return {
                "trial_id": trial_number,  # Will be updated after DB insert
                "trial_number": trial_number,
                "folder_path": folder_path,
                "folder_name": folder_name,
            }

    def complete_trial(
        self,
        trial_number: int,
        objectives: Dict[str, float],
        weighted_sum: Optional[float] = None,
        is_feasible: bool = True,
        metadata: Optional[Dict[str, Any]] = None
    ) -> int:
        """
        Complete a trial - logs to database and updates folder metadata.

        Args:
            trial_number: Trial number from new_trial()
            objectives: Objective values from FEA
            weighted_sum: Combined objective for ranking
            is_feasible: Whether constraints are satisfied
            metadata: Additional info (solve_time, prediction_error, etc.)

        Returns:
            Database trial_id
        """
        folder_path = self.iterations_dir / f"trial_{trial_number:04d}"

        # Load existing metadata
        meta_file = folder_path / "_meta.json"
        with open(meta_file, 'r') as f:
            meta = json.load(f)

        params = meta.get("params", {})

        # Update metadata
        meta["status"] = "COMPLETE"
        meta["datetime_complete"] = datetime.now().isoformat()
        meta["objectives"] = objectives
        meta["weighted_sum"] = weighted_sum
        meta["is_feasible"] = is_feasible
        if metadata:
            meta.update(metadata)

        # Save results.json
        results_file = folder_path / "results.json"
        with open(results_file, 'w') as f:
            json.dump({
                "objectives": objectives,
                "weighted_sum": weighted_sum,
                "is_feasible": is_feasible,
                "metadata": metadata or {}
            }, f, indent=2)

        # Update _meta.json
        with open(meta_file, 'w') as f:
            json.dump(meta, f, indent=2)

        # Log to database
        db_metadata = metadata or {}
        db_metadata["source"] = meta.get("source", "unknown")
        if "turbo_batch" in meta:
            db_metadata["turbo_batch"] = meta["turbo_batch"]
        if "predicted_ws" in meta:
            db_metadata["predicted_ws"] = meta["predicted_ws"]

        trial_id = self.db.log_trial(
            params=params,
            objectives=objectives,
            weighted_sum=weighted_sum,
            is_feasible=is_feasible,
            state="COMPLETE",
            datetime_start=meta.get("datetime_start"),
            datetime_complete=meta.get("datetime_complete"),
            metadata=db_metadata,
        )

        # Check if this is the new best
        best = self.db.get_best_trial()
        if best and best['trial_id'] == trial_id:
            self.db.mark_best(trial_id)
            meta["is_best"] = True
            with open(meta_file, 'w') as f:
                json.dump(meta, f, indent=2)

        return trial_id

    def fail_trial(self, trial_number: int, error: str):
        """Mark a trial as failed."""
        folder_path = self.iterations_dir / f"trial_{trial_number:04d}"
        meta_file = folder_path / "_meta.json"

        if meta_file.exists():
            with open(meta_file, 'r') as f:
                meta = json.load(f)
            meta["status"] = "FAIL"
            meta["error"] = error
            meta["datetime_complete"] = datetime.now().isoformat()
            with open(meta_file, 'w') as f:
                json.dump(meta, f, indent=2)

    def get_trial_folder(self, trial_number: int) -> Path:
        """Get folder path for a trial number."""
        return self.iterations_dir / f"trial_{trial_number:04d}"

    def get_all_trials(self) -> List[Dict[str, Any]]:
        """Get all completed trials from database."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        cursor.execute("""
            SELECT t.trial_id, t.number, tv.value
            FROM trials t
            JOIN trial_values tv ON t.trial_id = tv.trial_id
            WHERE t.state = 'COMPLETE'
            ORDER BY t.number
        """)

        trials = []
        for row in cursor.fetchall():
            trials.append({
                "trial_id": row[0],
                "number": row[1],
                "value": row[2]
            })

        conn.close()
        return trials

    def get_summary(self) -> Dict[str, Any]:
        """Get trial manager summary."""
        summary = self.db.get_summary()

        # Add folder count
        folders = list(self.iterations_dir.glob("trial_*"))
        summary["folder_count"] = len(folders)

        return summary

    def copy_model_files(self, source_dir: Path, trial_number: int) -> Path:
        """Copy NX model files to trial folder."""
        dest = self.get_trial_folder(trial_number)

        # Copy relevant files
        extensions = ['.prt', '.fem', '.sim', '.afm', '.op2', '.f06', '.dat']
        for ext in extensions:
            for src_file in source_dir.glob(f"*{ext}"):
                shutil.copy2(src_file, dest / src_file.name)

        return dest