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

"""Task executor for local simulation system (P4-M2).

Listens for tasks dispatched from Web backend, executes Motor-CAD
simulations, reports progress and results back to Web backend.

NOTE: All strings must be ASCII only. Chinese text uses \\uXXXX escapes.
"""
import json
import os
import sys
import time
import uuid
import threading
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any, Callable

try:
    import requests
except ImportError:
    requests = None


class TaskExecutor:
    """Executes simulation tasks dispatched from Web backend."""

    def __init__(
        self,
        web_base_url: str = "http://127.0.0.1:8000",
        task_dir: Optional[str] = None,
        on_progress: Optional[Callable] = None,
        on_complete: Optional[Callable] = None,
        on_error: Optional[Callable] = None,
    ):
        self.web_base_url = web_base_url.rstrip("/")
        self.task_dir = task_dir or os.path.join(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
            "output", "tasks"
        )
        os.makedirs(self.task_dir, exist_ok=True)
        self.on_progress = on_progress
        self.on_complete = on_complete
        self.on_error = on_error
        self._running = False
        self._current_task: Optional[Dict[str, Any]] = None
        self._stop_event = threading.Event()

    def fetch_pending_tasks(self) -> List[Dict[str, Any]]:
        """Fetch pending tasks from Web backend."""
        if requests is None:
            return self._scan_local_task_files()
        try:
            resp = requests.get(
                f"{self.web_base_url}/api/tasks",
                params={"status": "pending", "limit": 10},
                timeout=10,
            )
            if resp.status_code == 200:
                data = resp.json()
                return data.get("tasks", [])
        except Exception as e:
            if self.on_error:
                self.on_error(f"Fetch tasks failed: {str(e)}")
        return []

    def _scan_local_task_files(self) -> List[Dict[str, Any]]:
        """Scan local task directory for task files (fallback mode)."""
        tasks = []
        for fname in os.listdir(self.task_dir):
            if fname.endswith("_task.json"):
                fpath = os.path.join(self.task_dir, fname)
                try:
                    with open(fpath, "r", encoding="utf-8") as f:
                        task = json.load(f)
                    task["_local_file"] = fpath
                    tasks.append(task)
                except Exception:
                    continue
        return tasks

    def dispatch_task(self, task_id: str) -> bool:
        """Mark task as dispatched on Web backend."""
        if requests is None:
            return True
        try:
            resp = requests.post(
                f"{self.web_base_url}/api/tasks/{task_id}/dispatch",
                timeout=10,
            )
            return resp.status_code in (200, 201)
        except Exception as e:
            if self.on_error:
                self.on_error(f"Dispatch task {task_id} failed: {str(e)}")
            return False

    def report_progress(
        self,
        task_id: str,
        current_point: int,
        total_points: int,
        current_params: Optional[Dict[str, Any]] = None,
        elapsed_time: Optional[float] = None,
    ) -> bool:
        """Report simulation progress to Web backend."""
        if requests is None:
            if self.on_progress:
                self.on_progress(task_id, current_point, total_points)
            return True
        try:
            payload = {
                "current_point": current_point,
                "total_points": total_points,
                "current_params": current_params,
                "elapsed_time": elapsed_time,
            }
            resp = requests.post(
                f"{self.web_base_url}/api/tasks/{task_id}/progress",
                json=payload,
                timeout=10,
            )
            return resp.status_code == 200
        except Exception as e:
            if self.on_error:
                self.on_error(f"Report progress failed: {str(e)}")
            return False

    def report_results(
        self,
        task_id: str,
        results: List[Dict[str, Any]],
        metrics: Optional[Dict[str, Any]] = None,
        logs: Optional[str] = None,
        duration: Optional[float] = None,
        status: str = "completed",
    ) -> bool:
        """Report final results to Web backend."""
        if requests is None:
            if self.on_complete:
                self.on_complete(task_id, results)
            return True
        try:
            payload = {
                "results": results,
                "metrics": metrics,
                "logs": logs,
                "duration": duration,
                "status": status,
            }
            resp = requests.post(
                f"{self.web_base_url}/api/tasks/{task_id}/results",
                json=payload,
                timeout=30,
            )
            return resp.status_code == 200
        except Exception as e:
            if self.on_error:
                self.on_error(f"Report results failed: {str(e)}")
            return False

    def execute_task(self, task: Dict[str, Any]) -> None:
        """Execute a single simulation task.

        This is a template method. Override _run_simulation_point in
        subclasses to implement actual Motor-CAD simulation.
        """
        task_id = task.get("task_id", str(uuid.uuid4())[:8])
        parameters = task.get("parameters", [])
        total_points = len(parameters)
        results = []
        start_time = time.time()

        self._current_task = task
        self.dispatch_task(task_id)

        for idx, params in enumerate(parameters):
            if self._stop_event.is_set():
                break

            elapsed = time.time() - start_time
            self.report_progress(task_id, idx, total_points, params, elapsed)

            try:
                point_result = self._run_simulation_point(params, idx)
                point_result["point_index"] = idx
                point_result["params"] = params
                results.append(point_result)
            except Exception as e:
                results.append({
                    "point_index": idx,
                    "params": params,
                    "status": "failed",
                    "error": str(e),
                })
                if self.on_error:
                    self.on_error(f"Point {idx} failed: {str(e)}")

        duration = time.time() - start_time
        metrics = self._compute_metrics(results)
        status = "completed" if not self._stop_event.is_set() else "cancelled"

        self.report_results(task_id, results, metrics, None, duration, status)
        self.report_progress(task_id, total_points, total_points, None, duration)

        self._current_task = None
        if self.on_complete:
            self.on_complete(task_id, results, metrics)

    def _run_simulation_point(self, params: Dict[str, Any], index: int) -> Dict[str, Any]:
        """Run a single simulation point. Override in subclass.

        Template implementation returns mock data.
        """
        import random
        rng = random.Random(index + hash(json.dumps(params, sort_keys=True)) % 10000)
        airgap = params.get("airgap_mm", 1.0)
        current = params.get("current_a", 15.0)
        return {
            "tavg_nm": round(current * 2.5 / (airgap ** 0.5) + rng.gauss(0, 0.3), 4),
            "efficiency_pct": round(88 + rng.gauss(0, 2), 2),
            "total_losses_w": round(50 + rng.gauss(0, 10), 2),
            "winding_temp_c": round(90 + rng.gauss(0, 10), 1),
            "status": "ok",
        }

    def _compute_metrics(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Compute aggregated metrics from results."""
        ok_results = [r for r in results if r.get("status") == "ok"]
        if not ok_results:
            return {"total_points": len(results), "successful_points": 0}
        metrics = {
            "total_points": len(results),
            "successful_points": len(ok_results),
        }
        for key in ["tavg_nm", "efficiency_pct", "total_losses_w", "winding_temp_c"]:
            values = [r[key] for r in ok_results if key in r]
            if values:
                metrics[f"{key}_min"] = min(values)
                metrics[f"{key}_max"] = max(values)
                metrics[f"{key}_mean"] = round(sum(values) / len(values), 4)
        return metrics

    def start_polling(self, interval: int = 5) -> threading.Thread:
        """Start background thread to poll for and execute tasks."""
        self._running = True
        self._stop_event.clear()

        def poll_loop():
            while self._running and not self._stop_event.is_set():
                try:
                    tasks = self.fetch_pending_tasks()
                    for task in tasks:
                        if self._stop_event.is_set():
                            break
                        self.execute_task(task)
                except Exception as e:
                    if self.on_error:
                        self.on_error(f"Poll loop error: {str(e)}")
                self._stop_event.wait(interval)

        thread = threading.Thread(target=poll_loop, daemon=True)
        thread.start()
        return thread

    def stop(self):
        """Stop the executor."""
        self._running = False
        self._stop_event.set()


class MotorCADTaskExecutor(TaskExecutor):
    """Task executor that uses actual Motor-CAD for simulation.

    Overrides _run_simulation_point to call Motor-CAD via pymotorcad.
    Falls back to mock data if Motor-CAD is not available.
    """

    def __init__(self, *args, model_path: Optional[str] = None, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_path = model_path
        self._mc = None

    def _run_simulation_point(self, params: Dict[str, Any], index: int) -> Dict[str, Any]:
        """Run simulation using Motor-CAD. Falls back to mock if unavailable."""
        try:
            return self._run_motorcad(params, index)
        except Exception as e:
            if self.on_error:
                self.on_error(f"MotorCAD failed for point {index}, using mock: {str(e)}")
            return super()._run_simulation_point(params, index)

    def _run_motorcad(self, params: Dict[str, Any], index: int) -> Dict[str, Any]:
        """Actual Motor-CAD simulation. Requires pymotorcad and Motor-CAD."""
        try:
            from ansys.motorcad.core import MotorCAD
        except ImportError:
            raise RuntimeError("pymotorcad not installed")

        if self._mc is None:
            self._mc = MotorCAD()
            if self.model_path and os.path.exists(self.model_path):
                self._mc.load_from_file(self.model_path)

        # Write parameters
        for key, value in params.items():
            try:
                self._mc.set_variable(key, value)
                applied = float(self._mc.get_variable(key))
                if abs(applied - value) > 1e-6:
                    raise RuntimeError(f"Variable {key} write mismatch: {applied} != {value}")
            except Exception:
                pass

        # Run simulation
        self._mc.do_magnetic_calculation()

        # Read results
        result = {
            "tavg_nm": float(self._mc.get_variable("Torque_Avg")),
            "efficiency_pct": float(self._mc.get_variable("Efficiency")),
            "total_losses_w": float(self._mc.get_variable("Total_Losses")),
            "status": "ok",
        }
        return result


if __name__ == "__main__":
    # Standalone test: run executor with mock data
    executor = TaskExecutor(
        web_base_url=os.environ.get("WEB_BASE_URL", "http://127.0.0.1:8000"),
        on_progress=lambda tid, cur, tot: print(f"[{tid}] Progress: {cur}/{tot}"),
        on_complete=lambda tid, res, met: print(f"[{tid}] Complete: {len(res)} points"),
        on_error=lambda msg: print(f"ERROR: {msg}"),
    )
    print("Task executor started. Press Ctrl+C to stop.")
    try:
        thread = executor.start_polling(interval=5)
        while thread.is_alive():
            time.sleep(1)
    except KeyboardInterrupt:
        executor.stop()
        print("Executor stopped.")