Atomizer/atomizer.py

#!/usr/bin/env python
"""
Atomizer CLI - Neural-Accelerated Structural Optimization

One-command interface for the complete Atomizer workflow:
- Create studies from templates
- Run FEA optimizations with auto training data export
- Auto-train neural networks when data threshold is reached
- Run neural-accelerated optimization (2200x faster!)

Usage:
    python atomizer.py neural-optimize --study my_study --trials 500
    python atomizer.py create-study --template beam_stiffness --name my_beam
    python atomizer.py status --study my_study
    python atomizer.py train --study my_study --epochs 100

The neural-optimize command is the main entry point - it handles the complete
workflow automatically:
1. Runs FEA optimization with training data export
2. Triggers neural network training when enough data is collected
3. Switches to neural-accelerated mode for remaining trials
4. Detects model drift and retrains as needed
"""

import argparse
import json
import logging
import sys
import time
from pathlib import Path
from typing import Optional

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

from optimization_engine.processors.surrogates.auto_trainer import AutoTrainer, check_training_status
from optimization_engine.config.template_loader import (
    create_study_from_template,
    list_templates,
    get_template
)
from optimization_engine.validators.study_validator import (
    validate_study,
    list_studies,
    quick_check
)


def setup_logging(verbose: bool = False) -> None:
    """Configure logging."""
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(
        level=level,
        format='%(asctime)s [%(levelname)s] %(message)s',
        datefmt='%H:%M:%S'
    )


def cmd_neural_optimize(args) -> int:
    """
    Run neural-accelerated optimization.

    This is the main workflow that:
    1. Validates study setup
    2. Runs FEA exploration with training data export
    3. Auto-trains neural model when threshold reached
    4. Runs remaining trials with neural acceleration
    """
    print("=" * 60)
    print("ATOMIZER NEURAL-ACCELERATED OPTIMIZATION")
    print("=" * 60)
    print(f"Study: {args.study}")
    print(f"Total trials: {args.trials}")
    print(f"Auto-train threshold: {args.min_points} points")
    print(f"Retrain every: {args.retrain_every} new points")
    print("=" * 60)

    # Validate study
    print("\n[1/5] Validating study setup...")
    validation = validate_study(args.study)

    if not validation.is_ready_to_run:
        print(f"\nStudy validation failed:")
        print(validation)
        return 1

    print(f"  Study is ready to run")
    print(f"  Design variables: {validation.summary.get('design_variables', 0)}")
    print(f"  Objectives: {validation.summary.get('objectives', 0)}")

    # Initialize auto-trainer
    print("\n[2/5] Initializing auto-trainer...")
    trainer = AutoTrainer(
        study_name=args.study,
        min_points=args.min_points,
        epochs=args.epochs,
        retrain_threshold=args.retrain_every
    )

    status = trainer.get_status()
    print(f"  Current data points: {status['total_points']}")
    print(f"  Model version: v{status['model_version']}")

    # Determine workflow phase
    has_trained_model = status['model_version'] > 0
    current_points = status['total_points']

    if has_trained_model and current_points >= args.min_points:
        print("\n[3/5] Neural model available - starting neural-accelerated optimization...")
        return _run_neural_phase(args, trainer)
    else:
        print("\n[3/5] Building training dataset with FEA exploration...")
        return _run_exploration_phase(args, trainer)


def _run_exploration_phase(args, trainer: AutoTrainer) -> int:
    """Run FEA exploration to build training dataset."""
    study_dir = PROJECT_ROOT / "studies" / args.study
    run_script = study_dir / "run_optimization.py"

    if not run_script.exists():
        print(f"Error: run_optimization.py not found in {study_dir}")
        return 1

    # Calculate how many FEA trials we need
    current_points = trainer.count_training_points()
    needed_for_training = args.min_points - current_points

    if needed_for_training > 0:
        fea_trials = min(needed_for_training + 10, args.trials)  # Extra buffer
        print(f"\n  Running {fea_trials} FEA trials to build training data...")
        print(f"  (Need {needed_for_training} more points for neural training)")
    else:
        fea_trials = args.trials
        print(f"\n  Running {fea_trials} FEA trials...")

    # Run FEA optimization
    import subprocess

    cmd = [
        sys.executable,
        str(run_script),
        "--trials", str(fea_trials)
    ]

    if args.resume:
        cmd.append("--resume")

    print(f"\n[4/5] Executing: {' '.join(cmd)}")
    print("-" * 60)

    start_time = time.time()
    result = subprocess.run(cmd, cwd=str(study_dir))
    elapsed = time.time() - start_time

    print("-" * 60)
    print(f"FEA optimization completed in {elapsed/60:.1f} minutes")

    # Check if we can now train
    print("\n[5/5] Checking training data...")
    if trainer.should_train():
        print("  Threshold reached! Training neural model...")
        model_path = trainer.train()
        if model_path:
            print(f"  Neural model trained: {model_path}")
            print(f"\n  Re-run with --resume to continue with neural acceleration!")
        else:
            print("  Training failed - check logs")
    else:
        status = trainer.get_status()
        remaining = args.min_points - status['total_points']
        print(f"  {status['total_points']} points collected")
        print(f"  Need {remaining} more for neural training")

    return result.returncode


def _run_neural_phase(args, trainer: AutoTrainer) -> int:
    """Run neural-accelerated optimization."""
    study_dir = PROJECT_ROOT / "studies" / args.study
    run_script = study_dir / "run_optimization.py"

    if not run_script.exists():
        print(f"Error: run_optimization.py not found in {study_dir}")
        return 1

    # Run with neural acceleration
    import subprocess

    cmd = [
        sys.executable,
        str(run_script),
        "--trials", str(args.trials),
        "--enable-nn"
    ]

    if args.resume:
        cmd.append("--resume")

    print(f"\n[4/5] Executing: {' '.join(cmd)}")
    print("-" * 60)

    start_time = time.time()
    result = subprocess.run(cmd, cwd=str(study_dir))
    elapsed = time.time() - start_time

    print("-" * 60)
    print(f"Neural optimization completed in {elapsed/60:.1f} minutes")

    # Check for retraining
    print("\n[5/5] Checking if retraining needed...")
    if trainer.should_train():
        print("  New data accumulated - triggering retraining...")
        model_path = trainer.train()
        if model_path:
            print(f"  New model version: {model_path}")
    else:
        status = trainer.get_status()
        print(f"  {status['new_points_since_training']} new points since last training")
        print(f"  (Retrain threshold: {args.retrain_every})")

    return result.returncode


def cmd_create_study(args) -> int:
    """Create a new study from template."""
    print(f"Creating study '{args.name}' from template '{args.template}'...")

    try:
        study_path = create_study_from_template(
            template_name=args.template,
            study_name=args.name
        )
        print(f"\nSuccess! Study created at: {study_path}")
        return 0
    except FileNotFoundError as e:
        print(f"Error: {e}")
        return 1
    except FileExistsError as e:
        print(f"Error: {e}")
        return 1


def cmd_list_templates(args) -> int:
    """List available templates."""
    templates = list_templates()

    if not templates:
        print("No templates found in templates/")
        return 1

    print("\nAvailable Templates:")
    print("=" * 60)

    for t in templates:
        print(f"\n{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("\n" + "=" * 60)
    print("Use: atomizer create-study --template <name> --name <study_name>")

    return 0


def cmd_status(args) -> int:
    """Show study and training status."""
    if args.study:
        # Show specific study status
        print(f"\n=== Study: {args.study} ===\n")

        # Validation status
        validation = validate_study(args.study)
        print("VALIDATION STATUS")
        print("-" * 40)
        print(f"  Status: {validation.status.value}")
        print(f"  Ready to run: {validation.is_ready_to_run}")

        for key, value in validation.summary.items():
            print(f"  {key}: {value}")

        # Training status
        print("\nTRAINING DATA STATUS")
        print("-" * 40)
        status = check_training_status(args.study)
        print(f"  Data points: {status['total_points']}")
        print(f"  New since training: {status['new_points_since_training']}")
        print(f"  Model version: v{status['model_version']}")
        print(f"  Should train: {status['should_train']}")

        if status['latest_model']:
            print(f"  Latest model: {status['latest_model']}")

    else:
        # List all studies
        print("\nAll Studies:")
        print("=" * 60)

        studies = list_studies()
        if not studies:
            print("  No studies found in studies/")
            return 0

        for study in studies:
            icon = "[OK]" if study["is_ready"] else "[!]"
            trials_info = f"{study['trials']} trials" if study['trials'] > 0 else "no trials"
            pareto_info = f", {study['pareto']} Pareto" if study['pareto'] > 0 else ""
            print(f"  {icon} {study['name']}")
            print(f"      Status: {study['status']} ({trials_info}{pareto_info})")

    return 0


def cmd_train(args) -> int:
    """Trigger neural network training."""
    print(f"Training neural model for study: {args.study}")

    trainer = AutoTrainer(
        study_name=args.study,
        min_points=args.min_points,
        epochs=args.epochs
    )

    status = trainer.get_status()
    print(f"\nCurrent status:")
    print(f"  Data points: {status['total_points']}")
    print(f"  Min threshold: {args.min_points}")

    if args.force or trainer.should_train():
        if args.force and status['total_points'] < args.min_points:
            print(f"\nWarning: Force training with {status['total_points']} points (< {args.min_points})")

        print("\nStarting training...")
        model_path = trainer.train()

        if model_path:
            print(f"\nSuccess! Model saved to: {model_path}")
            return 0
        else:
            print("\nTraining failed - check logs")
            return 1
    else:
        needed = args.min_points - status['total_points']
        print(f"\nNot enough data for training. Need {needed} more points.")
        print("Use --force to train anyway.")
        return 1


def cmd_validate(args) -> int:
    """Validate study setup."""
    validation = validate_study(args.study)
    print(validation)
    return 0 if validation.is_ready_to_run else 1


def main():
    parser = argparse.ArgumentParser(
        description="Atomizer - Neural-Accelerated Structural Optimization",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Complete neural workflow
  python atomizer.py neural-optimize --study my_study --trials 500

  # Create study from template
  python atomizer.py create-study --template beam_stiffness --name my_beam

  # Check status
  python atomizer.py status --study my_study

  # Manual training
  python atomizer.py train --study my_study --epochs 100
        """
    )

    parser.add_argument("-v", "--verbose", action="store_true", help="Verbose output")

    subparsers = parser.add_subparsers(dest="command", help="Commands")

    # neural-optimize command
    neural_parser = subparsers.add_parser(
        "neural-optimize",
        help="Run neural-accelerated optimization (main workflow)"
    )
    neural_parser.add_argument("--study", "-s", required=True, help="Study name")
    neural_parser.add_argument("--trials", "-n", type=int, default=500, help="Total trials")
    neural_parser.add_argument("--min-points", type=int, default=50, help="Min points for training")
    neural_parser.add_argument("--retrain-every", type=int, default=50, help="Retrain after N new points")
    neural_parser.add_argument("--epochs", type=int, default=100, help="Training epochs")
    neural_parser.add_argument("--resume", action="store_true", help="Resume existing study")

    # create-study command
    create_parser = subparsers.add_parser("create-study", 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")

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

    # status command
    status_parser = subparsers.add_parser("status", help="Show status")
    status_parser.add_argument("--study", "-s", help="Study name (omit for all)")

    # train command
    train_parser = subparsers.add_parser("train", help="Train neural model")
    train_parser.add_argument("--study", "-s", required=True, help="Study name")
    train_parser.add_argument("--epochs", type=int, default=100, help="Training epochs")
    train_parser.add_argument("--min-points", type=int, default=50, help="Min points threshold")
    train_parser.add_argument("--force", action="store_true", help="Force training")

    # validate command
    validate_parser = subparsers.add_parser("validate", help="Validate study setup")
    validate_parser.add_argument("--study", "-s", required=True, help="Study name")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        return 0

    setup_logging(args.verbose)

    # Dispatch to command handler
    commands = {
        "neural-optimize": cmd_neural_optimize,
        "create-study": cmd_create_study,
        "list-templates": cmd_list_templates,
        "status": cmd_status,
        "train": cmd_train,
        "validate": cmd_validate
    }

    handler = commands.get(args.command)
    if handler:
        return handler(args)
    else:
        parser.print_help()
        return 1


if __name__ == "__main__":
    sys.exit(main())