Atomizer/optimization_engine/template_loader.py

"""
Template Loader for Atomizer Optimization Studies

Creates new studies from templates with automatic folder structure creation.

Usage:
    from optimization_engine.template_loader import create_study_from_template, list_templates

    # List available templates
    templates = list_templates()

    # Create a new study from template
    create_study_from_template(
        template_name="beam_stiffness_optimization",
        study_name="my_beam_study"
    )
"""

import json
import shutil
from pathlib import Path
from typing import Dict, Any, List, Optional
from datetime import datetime


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


def list_templates() -> List[Dict[str, Any]]:
    """
    List all available templates.

    Returns:
        List of template metadata dictionaries
    """
    templates = []

    if not TEMPLATES_DIR.exists():
        return templates

    for template_file in TEMPLATES_DIR.glob("*.json"):
        try:
            with open(template_file, 'r') as f:
                config = json.load(f)

            template_info = config.get("template_info", {})
            templates.append({
                "name": template_file.stem,
                "description": config.get("description", "No description"),
                "category": template_info.get("category", "general"),
                "analysis_type": template_info.get("analysis_type", "unknown"),
                "objectives": len(config.get("objectives", [])),
                "design_variables": len(config.get("design_variables", [])),
                "path": str(template_file)
            })
        except Exception as e:
            print(f"Warning: Could not load template {template_file}: {e}")

    return templates


def get_template(template_name: str) -> Optional[Dict[str, Any]]:
    """
    Load a template by name.

    Args:
        template_name: Name of the template (without .json extension)

    Returns:
        Template configuration dictionary or None if not found
    """
    template_path = TEMPLATES_DIR / f"{template_name}.json"

    if not template_path.exists():
        # Try with .json extension already included
        template_path = TEMPLATES_DIR / template_name
        if not template_path.exists():
            return None

    with open(template_path, 'r') as f:
        return json.load(f)


def create_study_from_template(
    template_name: str,
    study_name: str,
    studies_dir: Optional[Path] = None,
    overrides: Optional[Dict[str, Any]] = None
) -> Path:
    """
    Create a new study from a template.

    Args:
        template_name: Name of the template to use
        study_name: Name for the new study
        studies_dir: Base directory for studies (default: studies/)
        overrides: Dictionary of config values to override

    Returns:
        Path to the created study directory

    Raises:
        FileNotFoundError: If template doesn't exist
        FileExistsError: If study already exists
    """
    if studies_dir is None:
        studies_dir = STUDIES_DIR

    studies_dir = Path(studies_dir)

    # Load template
    template = get_template(template_name)
    if template is None:
        available = [t["name"] for t in list_templates()]
        raise FileNotFoundError(
            f"Template '{template_name}' not found. "
            f"Available templates: {available}"
        )

    # Check if study already exists
    study_path = studies_dir / study_name
    if study_path.exists():
        raise FileExistsError(
            f"Study '{study_name}' already exists at {study_path}. "
            "Choose a different name or delete the existing study."
        )

    # Create study directory structure
    setup_dir = study_path / "1_setup"
    model_dir = setup_dir / "model"
    results_dir = study_path / "2_results"

    setup_dir.mkdir(parents=True)
    model_dir.mkdir()
    results_dir.mkdir()

    # Customize template for this study
    config = template.copy()
    config["study_name"] = study_name
    config["created_from_template"] = template_name
    config["created_at"] = datetime.now().isoformat()

    # Update training data export path
    if "training_data_export" in config:
        export_dir = config["training_data_export"].get("export_dir", "")
        if "${study_name}" in export_dir:
            config["training_data_export"]["export_dir"] = export_dir.replace(
                "${study_name}", study_name
            )

    # Apply overrides
    if overrides:
        _deep_update(config, overrides)

    # Write configuration
    config_path = setup_dir / "optimization_config.json"
    with open(config_path, 'w') as f:
        json.dump(config, f, indent=2)

    # Create run_optimization.py
    run_script_content = _generate_run_script(study_name, config)
    run_script_path = study_path / "run_optimization.py"
    with open(run_script_path, 'w') as f:
        f.write(run_script_content)

    # Create README.md
    readme_content = _generate_study_readme(study_name, config, template_name)
    readme_path = study_path / "README.md"
    with open(readme_path, 'w') as f:
        f.write(readme_content)

    print(f"Created study '{study_name}' from template '{template_name}'")
    print(f"  Location: {study_path}")
    print(f"  Config: {config_path}")
    print(f"\nNext steps:")
    print(f"  1. Add your NX model files to: {model_dir}")
    print(f"  2. Update design variable bounds in optimization_config.json")
    print(f"  3. Run: python {run_script_path} --trials 50")

    return study_path


def _deep_update(base: Dict, updates: Dict) -> Dict:
    """Recursively update a dictionary."""
    for key, value in updates.items():
        if key in base and isinstance(base[key], dict) and isinstance(value, dict):
            _deep_update(base[key], value)
        else:
            base[key] = value
    return base


def _generate_run_script(study_name: str, config: Dict[str, Any]) -> str:
    """Generate the run_optimization.py script for a study."""
    return f'''"""
Optimization Runner for {study_name}

Auto-generated from template: {config.get('created_from_template', 'unknown')}
Created: {config.get('created_at', 'unknown')}

Usage:
    python run_optimization.py --trials 50
    python run_optimization.py --trials 25 --resume
    python run_optimization.py --trials 100 --enable-nn
"""

import sys
import argparse
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from optimization_engine.study_runner import run_study


def main():
    parser = argparse.ArgumentParser(description="{config.get('description', study_name)}")
    parser.add_argument('--trials', type=int, default=30, help='Number of trials to run')
    parser.add_argument('--resume', action='store_true', help='Resume existing study')
    parser.add_argument('--enable-nn', action='store_true', help='Enable neural network acceleration')
    parser.add_argument('--validate-only', action='store_true', help='Only validate setup, do not run')

    args = parser.parse_args()

    study_dir = Path(__file__).parent
    config_path = study_dir / "1_setup" / "optimization_config.json"

    if args.validate_only:
        from optimization_engine.validators import validate_study
        result = validate_study("{study_name}")
        print(result)
        return

    run_study(
        config_path=config_path,
        n_trials=args.trials,
        resume=args.resume,
        enable_neural=args.enable_nn
    )


if __name__ == "__main__":
    main()
'''


def _generate_study_readme(study_name: str, config: Dict[str, Any], template_name: str) -> str:
    """Generate a README.md for the study."""
    objectives = config.get("objectives", [])
    design_vars = config.get("design_variables", [])
    constraints = config.get("constraints", [])

    obj_list = "\n".join([f"- **{o.get('name', 'unnamed')}**: {o.get('goal', 'minimize')} - {o.get('description', '')}" for o in objectives])
    dv_list = "\n".join([f"- **{d.get('parameter', 'unnamed')}**: [{d.get('bounds', [0, 1])[0]}, {d.get('bounds', [0, 1])[1]}] - {d.get('description', '')}" for d in design_vars])
    const_list = "\n".join([f"- **{c.get('name', 'unnamed')}**: {c.get('type', 'less_than')} {c.get('threshold', 0)} - {c.get('description', '')}" for c in constraints])

    return f'''# {study_name}

{config.get('description', 'Optimization study')}

**Template**: {template_name}
**Created**: {config.get('created_at', 'unknown')}

## Engineering Context

{config.get('engineering_context', 'No context provided')}

## Objectives

{obj_list if obj_list else 'None defined'}

## Design Variables

{dv_list if dv_list else 'None defined'}

## Constraints

{const_list if const_list else 'None defined'}

## Setup Instructions

1. **Add NX Model Files**

   Copy your NX part (.prt), simulation (.sim), and FEM (.fem) files to:
   ```
   1_setup/model/
   ```

2. **Configure Design Variables**

   Edit `1_setup/optimization_config.json`:
   - Ensure `design_variables[].parameter` matches your NX expression names
   - Adjust bounds to your design space

3. **Validate Setup**

   ```bash
   python run_optimization.py --validate-only
   ```

## Running the Optimization

### Basic Run
```bash
python run_optimization.py --trials 50
```

### Resume Interrupted Run
```bash
python run_optimization.py --trials 25 --resume
```

### With Neural Network Acceleration
```bash
python run_optimization.py --trials 100 --enable-nn
```

## Results

After optimization, results are saved in `2_results/`:
- `study.db` - Optuna database with all trials
- `history.json` - Trial history
- `optimization_summary.json` - Summary with best parameters

## Visualization

View results with Optuna Dashboard:
```bash
optuna-dashboard sqlite:///2_results/study.db
```

Or generate a report:
```bash
python -m optimization_engine.generate_report {study_name}
```
'''


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Atomizer Template Loader")
    subparsers = parser.add_subparsers(dest="command", help="Commands")

    # List templates
    list_parser = subparsers.add_parser("list", help="List available templates")

    # Create study
    create_parser = subparsers.add_parser("create", help="Create study from template")
    create_parser.add_argument("--template", "-t", required=True, help="Template name")
    create_parser.add_argument("--name", "-n", required=True, help="Study name")

    args = parser.parse_args()

    if args.command == "list":
        templates = list_templates()
        if not templates:
            print("No templates found in templates/")
        else:
            print("Available templates:")
            print("-" * 60)
            for t in templates:
                print(f"  {t['name']}")
                print(f"    {t['description']}")
                print(f"    Category: {t['category']} | Analysis: {t['analysis_type']}")
                print(f"    Design vars: {t['design_variables']} | Objectives: {t['objectives']}")
                print()

    elif args.command == "create":
        try:
            study_path = create_study_from_template(
                template_name=args.template,
                study_name=args.name
            )
        except (FileNotFoundError, FileExistsError) as e:
            print(f"Error: {e}")
            sys.exit(1)

    else:
        parser.print_help()