pcb-afm-simulation-system/scripts/run_scan.py

"""Command-line entry point for parameter scan.

Usage:
    python scripts/run_scan.py --config scan_config.json
    python scripts/run_scan.py --model models/MARS.mot --var Airgap --start 0.6 --stop 1.5 --step 0.3

All source is ASCII.
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

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

from src.solver_core import MotorCADSolver  # noqa: E402
from src.scan_engine import (  # noqa: E402
    ScanEngine,
    generate_cartesian_points,
    values_inclusive,
    estimate_total_time,
)

DEFAULT_MODEL = PROJECT_ROOT / "models" / "MARS-12S10P_SSSR_D76-C150_V5.0-0819.mot"
DEFAULT_OUTPUT = PROJECT_ROOT / "output"


def _log(text: str) -> None:
    from datetime import datetime
    stamp = datetime.now().strftime("%H:%M:%S")
    print(f"[{stamp}] {text}", flush=True)


def git_preflight(project_root: Path) -> tuple[bool, str]:
    """Check that git repo exists, HEAD is valid, and tracked files are clean.
    Returns (ok, commit_or_error)."""
    try:
        safe = f"safe.directory={project_root.as_posix()}"
        base = ["git", "-c", safe]
        commit = subprocess.check_output(
            base + ["rev-parse", "--short", "HEAD"],
            cwd=project_root, text=True, stderr=subprocess.STDOUT,
        ).strip()
        dirty = subprocess.check_output(
            base + ["status", "--porcelain", "--untracked-files=no"],
            cwd=project_root, text=True, stderr=subprocess.STDOUT,
        ).strip()
        if dirty:
            return False, f"Tracked files have uncommitted changes:\n{dirty}"
        return True, commit
    except (OSError, subprocess.CalledProcessError) as exc:
        return False, f"Git preflight failed: {exc}"


def load_config(config_path: Path) -> dict:
    """Load scan configuration from JSON file."""
    with open(config_path, "r", encoding="utf-8") as f:
        return json.load(f)


def build_points_from_config(config: dict) -> tuple[list[dict], list[str]]:
    """Build scan points from config. Returns (points, var_names)."""
    variables = []
    var_names = []
    for v in config.get("variables", []):
        name = v["name"]
        var_names.append(name)
        if "values" in v:
            values = [float(x) for x in v["values"]]
        else:
            values = values_inclusive(
                float(v["start"]), float(v["stop"]), float(v["step"])
            )
        variables.append({"name": name, "display_name": v.get("display_name", name), "values": values})

    points = generate_cartesian_points(variables)
    return points, var_names


def build_points_from_args(args) -> tuple[list[dict], list[str]]:
    """Build scan points from CLI arguments (single variable only)."""
    values = values_inclusive(args.start, args.stop, args.step)
    variables = [{"name": args.var, "display_name": args.var, "values": values}]
    points = generate_cartesian_points(variables)
    return points, [args.var]


def main() -> int:
    parser = argparse.ArgumentParser(description="Motor-CAD parameter scan")
    parser.add_argument("--config", type=str, default=None,
                        help="Path to scan config JSON")
    parser.add_argument("--model", type=str, default=str(DEFAULT_MODEL),
                        help="Path to .mot model file")
    parser.add_argument("--output-dir", type=str, default=str(DEFAULT_OUTPUT),
                        help="Output directory")
    parser.add_argument("--scan-name", type=str, default="scan",
                        help="Name for this scan run")
    parser.add_argument("--var", type=str, default=None,
                        help="Variable name to scan (CLI mode, single var)")
    parser.add_argument("--start", type=float, default=None,
                        help="Start value (CLI mode)")
    parser.add_argument("--stop", type=float, default=None,
                        help="Stop value (CLI mode)")
    parser.add_argument("--step", type=float, default=None,
                        help="Step value (CLI mode)")
    parser.add_argument("--quit", action="store_true",
                        help="Close Motor-CAD after scan")
    parser.add_argument("--skip-git-check", action="store_true",
                        help="Skip git preflight check")
    args = parser.parse_args()

    # Git preflight
    if not args.skip_git_check:
        ok, msg = git_preflight(PROJECT_ROOT)
        if not ok:
            print(f"GIT PREFLIGHT FAILED: {msg}")
            print("Commit all changes first, or use --skip-git-check to bypass.")
            return 1
        _log(f"Git preflight passed (commit {msg})")

    # Build scan points
    if args.config:
        config = load_config(Path(args.config))
        model_path = config.get("model_path", args.model)
        scan_name = config.get("scan_name", args.scan_name)
        points, var_names = build_points_from_config(config)
    elif args.var and args.start is not None and args.stop is not None and args.step is not None:
        model_path = args.model
        scan_name = args.scan_name
        points, var_names = build_points_from_args(args)
    else:
        print("ERROR: Provide either --config or --var/--start/--stop/--step")
        return 1

    model_path = Path(model_path)
    if not model_path.exists():
        print(f"ERROR: Model file not found: {model_path}")
        return 1

    # Estimate time
    est = estimate_total_time(points)
    _log("=" * 60)
    _log(f"Scan: {scan_name}")
    _log(f"Model: {model_path.name}")
    _log(f"Variables: {var_names}")
    _log(f"Total points: {est['count']}")
    _log(f"Estimated time: {est['total_min']} min ({est['total_h']} h)")
    _log("=" * 60)
    _log("")

    # Run
    solver = MotorCADSolver(log_cb=_log)
    result = {}

    try:
        solver.connect()
        _log("")

        engine = ScanEngine(
            solver=solver,
            model_path=model_path,
            output_dir=args.output_dir,
            scan_name=scan_name,
            log_cb=_log,
        )
        result = engine.run(points, var_names=var_names)

        _log("")
        _log("=" * 60)
        _log("SCAN COMPLETE")
        _log("=" * 60)
        s = result["summary"]
        _log(f"OK: {s['ok']}  FAILED: {s['failed']}  SKIPPED: {s['skipped']}")
        _log(f"Results CSV: {result['csv_path']}")
        _log(f"Run directory: {result['run_dir']}")

    except KeyboardInterrupt:
        _log("Interrupted by user.")
    except Exception as exc:
        _log(f"FATAL: {type(exc).__name__}: {exc}")
        import traceback
        _log(traceback.format_exc())
        return 1
    finally:
        if args.quit:
            solver.disconnect()
        else:
            _log("")
            _log("Motor-CAD kept open for manual inspection.")
            _log("Use --quit flag to auto-close.")

    _log("")
    _log("Done.")
    return 0


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