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


# New UX System imports (lazy loaded to avoid import errors)
def get_intake_processor():
    from optimization_engine.intake import IntakeProcessor

    return IntakeProcessor


def get_validation_gate():
    from optimization_engine.validation import ValidationGate

    return ValidationGate


def get_report_generator():
    from optimization_engine.reporting.html_report import HTMLReportGenerator

    return HTMLReportGenerator


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


# ============================================================================
# NEW UX SYSTEM COMMANDS
# ============================================================================


def cmd_intake(args) -> int:
    """Process an intake folder into a study."""
    IntakeProcessor = get_intake_processor()

    # Determine inbox folder
    inbox_path = Path(args.folder)

    if not inbox_path.is_absolute():
        inbox_dir = PROJECT_ROOT / "studies" / "_inbox"
        if (inbox_dir / args.folder).exists():
            inbox_path = inbox_dir / args.folder
        elif (PROJECT_ROOT / "studies" / args.folder).exists():
            inbox_path = PROJECT_ROOT / "studies" / args.folder

    if not inbox_path.exists():
        print(f"Error: Folder not found: {inbox_path}")
        return 1

    print(f"Processing intake: {inbox_path}")
    print("=" * 60)

    def progress(message: str, percent: float):
        bar_width = 30
        filled = int(bar_width * percent)
        bar = "=" * filled + "-" * (bar_width - filled)
        print(f"\r[{bar}] {percent * 100:5.1f}% {message}", end="", flush=True)
        if percent >= 1.0:
            print()

    try:
        processor = IntakeProcessor(inbox_path, progress_callback=progress)
        context = processor.process(
            run_baseline=not args.skip_baseline,
            copy_files=True,
            run_introspection=True,
        )

        print("\n" + "=" * 60)
        print("INTAKE COMPLETE")
        print("=" * 60)

        summary = context.get_context_summary()
        print(f"\nStudy: {context.study_name}")
        print(f"Location: {processor.study_dir}")
        print(f"\nContext loaded:")
        print(f"  Model: {'Yes' if summary['has_model'] else 'No'}")
        print(f"  Introspection: {'Yes' if summary['has_introspection'] else 'No'}")
        print(f"  Baseline: {'Yes' if summary['has_baseline'] else 'No'}")
        print(
            f"  Expressions: {summary['num_expressions']} ({summary['num_dv_candidates']} candidates)"
        )

        if context.has_baseline:
            print(f"\nBaseline: {context.get_baseline_summary()}")

        if summary["warnings"]:
            print(f"\nWarnings:")
            for w in summary["warnings"]:
                print(f"  - {w}")

        print(f"\nNext: atomizer gate {context.study_name}")
        return 0

    except Exception as e:
        print(f"\nError: {e}")
        if args.verbose:
            import traceback

            traceback.print_exc()
        return 1


def cmd_gate(args) -> int:
    """Run validation gate before optimization."""
    ValidationGate = get_validation_gate()

    study_path = Path(args.study)
    if not study_path.is_absolute():
        study_path = PROJECT_ROOT / "studies" / args.study

    if not study_path.exists():
        print(f"Error: Study not found: {study_path}")
        return 1

    print(f"Validation Gate: {study_path.name}")
    print("=" * 60)

    def progress(message: str, percent: float):
        bar_width = 30
        filled = int(bar_width * percent)
        bar = "=" * filled + "-" * (bar_width - filled)
        print(f"\r[{bar}] {percent * 100:5.1f}% {message}", end="", flush=True)
        if percent >= 1.0:
            print()

    try:
        gate = ValidationGate(study_path, progress_callback=progress)
        result = gate.validate(
            run_test_trials=not args.skip_trials,
            n_test_trials=args.trials,
        )

        print("\n" + "=" * 60)
        if result.passed:
            print("VALIDATION PASSED")
        else:
            print("VALIDATION FAILED")
        print("=" * 60)

        # Show test trials
        if result.test_trials:
            print(
                f"\nTest Trials: {len([t for t in result.test_trials if t.success])}/{len(result.test_trials)} passed"
            )

            if result.results_vary:
                print("Results vary: Yes (mesh updating correctly)")
            else:
                print("Results vary: NO - MESH MAY NOT BE UPDATING!")

            # Results table
            print(f"\n{'Trial':<8} {'Status':<8} {'Time':<8}", end="")
            if result.test_trials and result.test_trials[0].objectives:
                for obj in list(result.test_trials[0].objectives.keys())[:3]:
                    print(f" {obj[:10]:<12}", end="")
            print()

            for trial in result.test_trials:
                status = "OK" if trial.success else "FAIL"
                print(
                    f"{trial.trial_number:<8} {status:<8} {trial.solve_time_seconds:<8.1f}", end=""
                )
                for val in list(trial.objectives.values())[:3]:
                    print(f" {val:<12.4f}", end="")
                print()

        # Runtime estimate
        if result.avg_solve_time:
            print(f"\nRuntime Estimate:")
            print(f"  Avg solve: {result.avg_solve_time:.1f}s")
            if result.estimated_total_runtime:
                print(f"  Total: {result.estimated_total_runtime / 3600:.1f}h")

        # Errors
        if result.errors:
            print(f"\nErrors:")
            for err in result.errors:
                print(f"  - {err}")

        if result.passed and args.approve:
            gate.approve()
            print(f"\nStudy approved for optimization!")
        elif result.passed:
            print(f"\nTo approve: atomizer gate {args.study} --approve")

        gate.save_result(result)
        return 0 if result.passed else 1

    except Exception as e:
        print(f"\nError: {e}")
        if args.verbose:
            import traceback

            traceback.print_exc()
        return 1


def cmd_finalize(args) -> int:
    """Generate final report for a study."""
    HTMLReportGenerator = get_report_generator()

    study_path = Path(args.study)
    if not study_path.is_absolute():
        study_path = PROJECT_ROOT / "studies" / args.study

    if not study_path.exists():
        print(f"Error: Study not found: {study_path}")
        return 1

    print(f"Generating report for: {study_path.name}")
    print("=" * 60)

    try:
        generator = HTMLReportGenerator(study_path)
        report_path = generator.generate(include_pdf=getattr(args, "pdf", False))

        print(f"\nReport generated successfully!")
        print(f"  HTML: {report_path}")
        print(f"  Data: {report_path.parent / 'data'}")

        if getattr(args, "open", False):
            import webbrowser

            webbrowser.open(str(report_path))
        else:
            print(f"\nOpen in browser: file://{report_path}")

        return 0

    except Exception as e:
        print(f"\nError: {e}")
        if args.verbose:
            import traceback

            traceback.print_exc()
        return 1


def cmd_list_studies(args) -> int:
    """List all studies and inbox items."""
    studies_dir = PROJECT_ROOT / "studies"

    print("Atomizer Studies")
    print("=" * 60)

    # Inbox items
    inbox_dir = studies_dir / "_inbox"
    if inbox_dir.exists():
        inbox_items = [d for d in inbox_dir.iterdir() if d.is_dir() and not d.name.startswith(".")]
        if inbox_items:
            print("\nPending Intake (_inbox/):")
            for item in sorted(inbox_items):
                has_config = (item / "intake.yaml").exists()
                has_model = bool(list(item.glob("**/*.sim")))
                status = []
                if has_config:
                    status.append("yaml")
                if has_model:
                    status.append("model")
                print(f"  {item.name:<30} [{', '.join(status) or 'empty'}]")

    # Active studies
    print("\nStudies:")
    for study_dir in sorted(studies_dir.iterdir()):
        if (
            study_dir.is_dir()
            and not study_dir.name.startswith("_")
            and not study_dir.name.startswith(".")
        ):
            has_spec = (study_dir / "atomizer_spec.json").exists() or (
                study_dir / "optimization_config.json"
            ).exists()
            has_db = any(study_dir.rglob("study.db"))
            has_approval = (study_dir / ".validation_approved").exists()

            status = []
            if has_spec:
                status.append("configured")
            if has_approval:
                status.append("approved")
            if has_db:
                status.append("has_data")

            print(f"  {study_dir.name:<30} [{', '.join(status) or 'new'}]")

    return 0


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")

    # ========================================================================
    # NEW UX SYSTEM COMMANDS
    # ========================================================================

    # intake command
    intake_parser = subparsers.add_parser("intake", help="Process an intake folder into a study")
    intake_parser.add_argument("folder", help="Path to intake folder")
    intake_parser.add_argument("--skip-baseline", action="store_true", help="Skip baseline solve")

    # gate command (validation gate)
    gate_parser = subparsers.add_parser("gate", help="Run validation gate with test trials")
    gate_parser.add_argument("study", help="Study name or path")
    gate_parser.add_argument("--skip-trials", action="store_true", help="Skip test trials")
    gate_parser.add_argument("--trials", type=int, default=3, help="Number of test trials")
    gate_parser.add_argument("--approve", action="store_true", help="Approve if validation passes")

    # list command
    list_studies_parser = subparsers.add_parser("list", help="List all studies and inbox items")

    # finalize command
    finalize_parser = subparsers.add_parser("finalize", help="Generate final HTML report")
    finalize_parser.add_argument("study", help="Study name or path")
    finalize_parser.add_argument("--pdf", action="store_true", help="Also generate PDF")
    finalize_parser.add_argument("--open", action="store_true", help="Open report in browser")

    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,
        # New UX commands
        "intake": cmd_intake,
        "gate": cmd_gate,
        "list": cmd_list_studies,
        "finalize": cmd_finalize,
    }

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


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