pcb-afm-simulation-system/src/scan_engine.py

"""Parameter scan engine for PCB axial flux motor simulation.

Supports single-parameter and multi-parameter Cartesian product scans
with per-point baseline reload, write-back verification, checkpoint
resume, and per-point result persistence (CSV + raw + manifest + log).

Built on top of MotorCADSolver from solver_core.

All source is ASCII.
"""

from __future__ import annotations

import csv
import json
import math
import time
import traceback
from datetime import datetime
from pathlib import Path
from typing import Callable

from .solver_core import (
    MotorCADSolver,
    METRIC_KEYS,
    METRIC_LABELS,
)


# ---------------------------------------------------------------------------
# Utility functions
# ---------------------------------------------------------------------------

def values_inclusive(start: float, stop: float, step: float) -> list[float]:
    """Generate values from start to stop inclusive, appending stop if
    it is not exactly reachable by integer steps."""
    if step <= 0 or stop < start:
        raise ValueError("step must be positive and stop must be >= start")
    count = int(math.floor((stop - start) / step + 1e-9))
    values = [round(start + i * step, 10) for i in range(count + 1)]
    if not math.isclose(values[-1], stop, abs_tol=1e-9):
        values.append(float(stop))
    return values


def generate_cartesian_points(variables: list[dict]) -> list[dict]:
    """Generate all parameter combinations via Cartesian product.

    Each variable dict: {"name": str, "display_name": str, "values": [float]}
    Returns list of {"index": int, "params": {var_name: value, ...}}
    """
    if not variables:
        return [{"index": 1, "params": {}}]

    # Build list of (name, values) pairs
    var_list = [(v["name"], v["values"]) for v in variables]

    # Recursive Cartesian product
    def _combine(idx: int, current: dict) -> list[dict]:
        if idx == len(var_list):
            return [{"params": dict(current)}]
        name, vals = var_list[idx]
        results = []
        for v in vals:
            current[name] = v
            results.extend(_combine(idx + 1, current))
        return results

    points = _combine(0, {})
    for i, p in enumerate(points, 1):
        p["index"] = i
    return points


def estimate_total_time(points: list[dict], per_point_s: float = 140.0) -> dict:
    """Estimate total scan time. Returns dict with count, per_point_s, total_s, total_h."""
    n = len(points)
    total_s = n * per_point_s
    return {
        "count": n,
        "per_point_s": per_point_s,
        "total_s": round(total_s, 1),
        "total_min": round(total_s / 60, 1),
        "total_h": round(total_s / 3600, 2),
    }


# ---------------------------------------------------------------------------
# Scan engine
# ---------------------------------------------------------------------------

class ScanEngine:
    """Runs a parameter scan using a connected MotorCADSolver.

    Per point: reload baseline model -> write params (with write-back
    verify) -> magnetic calculation -> export results -> extract metrics
    -> append CSV row (flushed immediately).

    Supports checkpoint resume: points already present in the output CSV
    with status OK are skipped.
    """

    def __init__(
        self,
        solver: MotorCADSolver,
        model_path: str | Path,
        output_dir: str | Path = "output",
        scan_name: str = "scan",
        log_cb: Callable[[str], None] | None = None,
        progress_cb: Callable[[int, int], None] | None = None,
        row_cb: Callable[[dict], None] | None = None,
    ):
        self.solver = solver
        self.model_path = str(Path(model_path).resolve())
        self.output_dir = Path(output_dir)
        self.scan_name = scan_name
        self.log_cb = log_cb
        self.progress_cb = progress_cb
        self.row_cb = row_cb
        self._cancel = False

    def _log(self, text: str) -> None:
        if self.log_cb:
            self.log_cb(text)

    def cancel(self) -> None:
        """Request cancellation after the current point finishes."""
        self._cancel = True

    def _csv_fields(self, var_names: list[str]) -> list[str]:
        """Build CSV field list: index, status, seconds, error, params, metrics."""
        return ["run_index", "status", "seconds", "error"] + var_names + METRIC_KEYS

    def _load_completed_indices(self, csv_path: Path) -> set[int]:
        """Load set of already-completed (OK) point indices from CSV."""
        completed = set()
        if not csv_path.exists():
            return completed
        try:
            with open(csv_path, "r", encoding="utf-8-sig", newline="") as f:
                reader = csv.DictReader(f)
                for row in reader:
                    if row.get("status") == "OK":
                        try:
                            completed.add(int(row["run_index"]))
                        except (ValueError, KeyError):
                            pass
        except Exception:
            pass
        return completed

    def run(
        self,
        points: list[dict],
        var_names: list[str] | None = None,
    ) -> dict:
        """Run the full scan.

        Args:
            points: list of {"index": int, "params": {var_name: value}}
            var_names: ordered list of variable names for CSV columns.
                       If None, extracted from first point's params keys.

        Returns:
            dict with run_dir, csv_path, log_path, manifest_path, summary.
        """
        if var_names is None:
            var_names = list(points[0]["params"].keys()) if points else []

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
        run_dir = self.output_dir / f"{timestamp}_{self.scan_name}"
        raw_dir = run_dir / "raw"
        run_dir.mkdir(parents=True, exist_ok=True)
        raw_dir.mkdir(parents=True, exist_ok=True)

        csv_path = run_dir / f"scan_results_{timestamp}.csv"
        log_path = run_dir / f"program_log_{timestamp}.log"
        manifest_path = run_dir / f"run_manifest_{timestamp}.json"

        csv_fields = self._csv_fields(var_names)

        # Checkpoint: load already completed indices
        completed = self._load_completed_indices(csv_path)
        if completed:
            self._log(f"Checkpoint: {len(completed)} points already completed, will skip")

        # Write manifest
        manifest = {
            "timestamp": timestamp,
            "model": self.model_path,
            "scan_name": self.scan_name,
            "total_points": len(points),
            "variable_names": var_names,
            "points": [
                {
                    "index": p["index"],
                    "params": p["params"],
                }
                for p in points
            ],
        }
        manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8")

        # Open log file
        log_file = open(log_path, "a", encoding="ascii", errors="backslashreplace")

        def _file_log(text: str) -> None:
            stamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
            line = f"{stamp} {text}"
            log_file.write(line + "\n")
            log_file.flush()
            self._log(line)

        _file_log(f"Scan started: {self.scan_name}")
        _file_log(f"Model: {self.model_path}")
        _file_log(f"Total points: {len(points)}")
        _file_log(f"Variables: {var_names}")
        _file_log(f"Run directory: {run_dir}")

        summary = {"ok": 0, "failed": 0, "skipped": 0, "results": []}

        try:
            # Open CSV for appending (checkpoint support)
            file_exists = csv_path.exists()
            with open(csv_path, "a", newline="", encoding="utf-8-sig") as csv_file:
                writer = csv.DictWriter(csv_file, fieldnames=csv_fields)
                if not file_exists:
                    writer.writeheader()
                    csv_file.flush()

                for point in points:
                    if self._cancel:
                        _file_log("Cancel requested; stopping before next point")
                        break

                    idx = point["index"]
                    params = point["params"]

                    # Checkpoint skip
                    if idx in completed:
                        summary["skipped"] += 1
                        _file_log(f"[{idx}/{len(points)}] SKIP (already completed)")
                        if self.progress_cb:
                            self.progress_cb(idx, len(points))
                        continue

                    param_str = ", ".join(f"{k}={v:g}" for k, v in params.items())
                    _file_log(f"[{idx}/{len(points)}] {param_str}")

                    # Run single point via solver
                    result = self.solver.run_single(
                        model_path=self.model_path,
                        params=params,
                        output_dir=raw_dir,
                        tag=f"pt{idx:04d}",
                    )

                    # Build CSV row
                    row = {field: "" for field in csv_fields}
                    row["run_index"] = idx
                    row["status"] = result["status"]
                    row["seconds"] = result["solve_time_s"]
                    row["error"] = result.get("error", "")
                    for k, v in params.items():
                        if k in row:
                            row[k] = v
                    for k, v in result.get("metrics", {}).items():
                        if k in row:
                            row[k] = v

                    writer.writerow(row)
                    csv_file.flush()

                    if result["status"] == "OK":
                        summary["ok"] += 1
                    else:
                        summary["failed"] += 1
                    summary["results"].append(row)

                    if self.row_cb:
                        self.row_cb(row)
                    if self.progress_cb:
                        self.progress_cb(idx, len(points))

                    ripple = row.get("ripple_pct", "")
                    tavg = row.get("tavg_nm", "")
                    eff = row.get("efficiency_pct", "")
                    _file_log(
                        f"[{idx}/{len(points)}] {row['status']} "
                        f"Tavg={tavg} ripple={ripple} eff={eff} "
                        f"seconds={row['seconds']}"
                    )

            _file_log(f"Scan ended. OK={summary['ok']} FAILED={summary['failed']} SKIPPED={summary['skipped']}")
            _file_log(f"Results CSV: {csv_path}")

        except Exception as exc:
            _file_log(f"FATAL: {type(exc).__name__}: {exc}")
            _file_log(traceback.format_exc())
            raise
        finally:
            log_file.close()

        return {
            "run_dir": str(run_dir),
            "csv_path": str(csv_path),
            "log_path": str(log_path),
            "manifest_path": str(manifest_path),
            "summary": summary,
        }