Atomizer/optimization_engine/templates/__init__.py

"""
Template Registry for Atomizer

Provides study templates for common optimization scenarios.
Used by Claude to quickly create new studies via wizard-driven workflow.
"""

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


REGISTRY_PATH = Path(__file__).parent / "registry.json"


def load_registry() -> Dict[str, Any]:
    """Load the template registry."""
    with open(REGISTRY_PATH, 'r') as f:
        return json.load(f)


def list_templates() -> List[Dict[str, Any]]:
    """List all available templates with summary info."""
    registry = load_registry()
    templates = []

    for t in registry["templates"]:
        templates.append({
            "id": t["id"],
            "name": t["name"],
            "description": t["description"],
            "category": t["category"],
            "n_objectives": len(t["objectives"]),
            "turbo_suitable": t.get("turbo_suitable", False),
            "example_study": t.get("example_study")
        })

    return templates


def get_template(template_id: str) -> Optional[Dict[str, Any]]:
    """Get a specific template by ID."""
    registry = load_registry()

    for t in registry["templates"]:
        if t["id"] == template_id:
            return t

    return None


def get_templates_by_category(category: str) -> List[Dict[str, Any]]:
    """Get all templates in a category."""
    registry = load_registry()

    return [t for t in registry["templates"] if t["category"] == category]


def list_categories() -> Dict[str, Dict[str, str]]:
    """List all template categories."""
    registry = load_registry()
    return registry.get("categories", {})


def get_extractor_info(extractor_id: str) -> Optional[Dict[str, Any]]:
    """Get information about a specific extractor."""
    registry = load_registry()
    return registry.get("extractors", {}).get(extractor_id)


def suggest_template(
    n_objectives: int = 1,
    physics_type: str = "structural",
    element_types: Optional[List[str]] = None
) -> Optional[Dict[str, Any]]:
    """
    Suggest a template based on problem characteristics.

    Args:
        n_objectives: Number of objectives (1 = single, 2+ = multi)
        physics_type: Type of physics (structural, dynamics, optics, multiphysics)
        element_types: List of element types in the mesh

    Returns:
        Best matching template or None
    """
    registry = load_registry()
    candidates = []

    for t in registry["templates"]:
        score = 0

        # Match number of objectives
        t_obj = len(t["objectives"])
        if n_objectives == 1 and t_obj == 1:
            score += 10
        elif n_objectives > 1 and t_obj > 1:
            score += 10

        # Match category
        if t["category"] == physics_type:
            score += 20

        # Match element types
        if element_types:
            t_elements = set(t.get("element_types", []))
            user_elements = set(element_types)
            if t_elements & user_elements:
                score += 15
            if "CQUAD4" in user_elements and "shell" in t["id"].lower():
                score += 10

        if score > 0:
            candidates.append((score, t))

    if not candidates:
        return None

    # Sort by score descending
    candidates.sort(key=lambda x: x[0], reverse=True)
    return candidates[0][1]


def format_template_summary(template: Dict[str, Any]) -> str:
    """Format a template as a human-readable summary."""
    lines = [
        f"**{template['name']}**",
        f"_{template['description']}_",
        "",
        f"**Category**: {template['category']}",
        f"**Solver**: {template.get('solver', 'SOL 101')}",
        "",
        "**Objectives**:"
    ]

    for obj in template["objectives"]:
        lines.append(f"  - {obj['name']} ({obj['direction']}) → Extractor {obj['extractor']}")

    lines.append("")
    lines.append("**Recommended Trials**:")
    trials = template.get("recommended_trials", {})
    for phase, count in trials.items():
        lines.append(f"  - {phase}: {count}")

    if template.get("turbo_suitable"):
        lines.append("")
        lines.append("✅ **Turbo Mode**: Suitable for neural acceleration")

    if template.get("notes"):
        lines.append("")
        lines.append(f"⚠️ **Note**: {template['notes']}")

    if template.get("example_study"):
        lines.append("")
        lines.append(f"📁 **Example**: studies/{template['example_study']}/")

    return "\n".join(lines)


def get_wizard_questions(template_id: str) -> List[Dict[str, Any]]:
    """Get wizard questions for a template."""
    template = get_template(template_id)
    if not template:
        return []
    return template.get("wizard_questions", [])


if __name__ == "__main__":
    # Demo: list all templates
    print("=== Atomizer Template Registry ===\n")

    for category_id, category in list_categories().items():
        print(f"{category['icon']} {category['name']}")
        print(f"   {category['description']}\n")

    print("\n=== Available Templates ===\n")

    for t in list_templates():
        status = "🚀" if t["turbo_suitable"] else "📊"
        print(f"{status} {t['name']} ({t['id']})")
        print(f"   {t['description']}")
        print(f"   Objectives: {t['n_objectives']} | Example: {t['example_study'] or 'N/A'}")
        print()