Update app.py

app.py

@@ -1,15 +1,18 @@
-import time
 import math
 import random
+import time
 import threading
 import collections
 from dataclasses import dataclass, asdict
+from pathlib import Path
 from typing import Optional, List, Dict, Any, Literal
 
-from fastapi import FastAPI
-from fastapi.responses import HTMLResponse, FileResponse
+from fastapi import FastAPI, Request
+from fastapi.responses import HTMLResponse, FileResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
+
+BASE_DIR = Path(__file__).resolve().parent
+INDEX_PATH = BASE_DIR / "index.html"
 
 app = FastAPI()
 app.add_middleware(
@@ -22,15 +25,15 @@ app.add_middleware(
 
 @dataclass
 class EngineConfig:
-    architecture: str = "additive"      # additive | multiplicative | affine | bilinear | gated
-    coeff_mode: str = "single_k"        # single_k | triple_k | per_edge_k
-    topology: str = "single_cell"       # single_cell | chain | mesh
-    dataset_family: str = "housing"     # housing | subtraction | multiplication | mixed | symbolic
-    mode: str = "training"              # training | inference
-    num_cells: int = 3
+    mode: Literal["training", "inference"] = "training"
+    architecture: Literal["additive", "multiplicative", "affine", "gated"] = "additive"
+    coeff_mode: Literal["single_k", "triple_k"] = "single_k"
+    topology: Literal["single_cell", "chain", "mesh"] = "single_cell"
+    dataset_family: Literal["housing", "subtraction", "multiplication", "symbolic", "mixed"] = "housing"
+    num_cells: int = 4
     learning_rate: float = 0.01
-    damping: float = 0.12
     coupling: float = 0.05
+    damping: float = 0.08
     batch_size: int = 24
     sample_noise: float = 0.0
 
@@ -49,7 +52,7 @@ class CellState:
     kc: float = 0.0
     prediction: float = 0.0
     error: float = 0.0
-    energy: float = 0.0
+    loss: float = 0.0
     force: float = 0.0
     anchored: bool = False
 
@@ -61,36 +64,37 @@ class SimEngine:
     def __init__(self):
         self.lock = threading.Lock()
         self.config = EngineConfig()
-
         self.running = False
         self.iteration = 0
         self.current_error = 0.0
         self.current_loss = 0.0
-        self.logs = []
-
+        self.last_action = "Idle"
+        self.logs: List[str] = []
         self.cells: List[CellState] = []
         self.batch_queue = collections.deque()
         self.current_sample: Optional[Dict[str, Any]] = None
         self.last_sample: Optional[Dict[str, Any]] = None
-
         self.loss_history = collections.deque(maxlen=120)
         self.error_history = collections.deque(maxlen=120)
-
         self.reset_state()
 
+    def add_log(self, msg: str):
+        self.logs.insert(0, f"[{self.iteration}] {msg}")
+        self.logs = self.logs[:20]
+        self.last_action = msg
+
     def reset_state(self):
-        with self.lock:
-            self.iteration = 0
-            self.current_error = 0.0
-            self.current_loss = 0.0
-            self.logs = []
-            self.batch_queue.clear()
-            self.current_sample = None
-            self.last_sample = None
-            self.loss_history.clear()
-            self.error_history.clear()
-            self._build_cells()
-            self.add_log("Engine reset.")
+        self.iteration = 0
+        self.current_error = 0.0
+        self.current_loss = 0.0
+        self.logs = []
+        self.batch_queue.clear()
+        self.current_sample = None
+        self.last_sample = None
+        self.loss_history.clear()
+        self.error_history.clear()
+        self._build_cells()
+        self.add_log("Engine reset.")
 
     def _build_cells(self):
         count = 1 if self.config.topology == "single_cell" else max(2, int(self.config.num_cells))
@@ -99,434 +103,349 @@ class SimEngine:
             self.cells.append(
                 CellState(
                     id=i,
-                    a=0.0,
-                    b=0.0,
-                    c=0.0,
-                    k=random.uniform(0.35, 1.25),
-                    ka=random.uniform(0.35, 1.25),
-                    kb=random.uniform(0.35, 1.25),
-                    kc=random.uniform(-0.25, 0.25),
+                    k=random.uniform(0.6, 1.4),
+                    ka=random.uniform(0.6, 1.4),
+                    kb=random.uniform(0.6, 1.4),
+                    kc=random.uniform(-0.3, 0.3),
                 )
             )
 
-    def add_log(self, msg: str):
-        stamp = f"[{self.iteration}] {msg}"
-        self.logs.insert(0, stamp)
-        if len(self.logs) > 20:
-            self.logs.pop()
-
     def configure(self, payload: Dict[str, Any]):
         with self.lock:
+            self.config.mode = payload.get("mode", self.config.mode)
             self.config.architecture = payload.get("architecture", self.config.architecture)
             self.config.coeff_mode = payload.get("coeff_mode", self.config.coeff_mode)
             self.config.topology = payload.get("topology", self.config.topology)
             self.config.dataset_family = payload.get("dataset_family", self.config.dataset_family)
-            self.config.mode = payload.get("mode", self.config.mode)
             self.config.num_cells = int(payload.get("num_cells", self.config.num_cells))
             self.config.learning_rate = float(payload.get("learning_rate", self.config.learning_rate))
-            self.config.damping = float(payload.get("damping", self.config.damping))
             self.config.coupling = float(payload.get("coupling", self.config.coupling))
+            self.config.damping = float(payload.get("damping", self.config.damping))
             self.config.batch_size = int(payload.get("batch_size", self.config.batch_size))
             self.config.sample_noise = float(payload.get("sample_noise", self.config.sample_noise))
 
             self.running = False
             self.reset_state()
             self.add_log(
-                f"Config applied: {self.config.architecture} | {self.config.coeff_mode} | "
-                f"{self.config.topology} | {self.config.dataset_family} | {self.config.mode}"
+                f"Config: {self.config.mode} | {self.config.architecture} | "
+                f"{self.config.coeff_mode} | {self.config.topology} | {self.config.dataset_family}"
             )
 
-    def _sample_housing(self):
-        a = random.uniform(2, 10)
-        b = random.uniform(2, 10)
-        c = (2.5 * a) + (1.2 * b) + random.uniform(-self.config.sample_noise, self.config.sample_noise)
-        return a, b, c, "housing_affine"
-
-    def _sample_subtraction(self):
-        a = random.uniform(2, 10)
-        b = random.uniform(2, 10)
-        c = (1.0 * a) + (-1.0 * b) + random.uniform(-self.config.sample_noise, self.config.sample_noise)
-        return a, b, c, "signed_subtraction"
-
-    def _sample_multiplication(self):
+    def _make_sample(self, family: str):
         a = random.uniform(2, 10)
         b = random.uniform(2, 10)
-        c = (a * b) + random.uniform(-self.config.sample_noise, self.config.sample_noise)
-        return a, b, c, "multiplicative"
-
-    def _sample_symbolic(self):
-        a = random.uniform(1, 12)
-        b = random.uniform(1, 12)
-        branch = random.choice(["affine", "signed_affine", "hybrid"])
-        if branch == "affine":
-            c = (1.7 * a) + (0.9 * b)
-        elif branch == "signed_affine":
-            c = (0.8 * a) + (-1.4 * b) + 2.0
-        else:
-            c = (a * 0.6) + (b * 0.4) + ((a * b) * 0.2)
-        c += random.uniform(-self.config.sample_noise, self.config.sample_noise)
-        return a, b, c, f"symbolic_{branch}"
 
-    def generate_sample(self, family: Optional[str] = None) -> Dict[str, Any]:
-        family = family or self.config.dataset_family
         if family == "housing":
-            a, b, c, label = self._sample_housing()
+            c = 2.5 * a + 1.2 * b
+            label = "affine_housing"
         elif family == "subtraction":
-            a, b, c, label = self._sample_subtraction()
+            c = 1.0 * a - 1.0 * b
+            label = "signed_subtraction"
         elif family == "multiplication":
-            a, b, c, label = self._sample_multiplication()
+            c = a * b
+            label = "product"
         elif family == "symbolic":
-            a, b, c, label = self._sample_symbolic()
+            mode = random.choice(["affine", "hybrid", "gated"])
+            if mode == "affine":
+                c = 1.7 * a + 0.9 * b + 1.0
+            elif mode == "hybrid":
+                c = 0.5 * a + 0.4 * b + 0.2 * a * b
+            else:
+                c = (a + b) / (1.0 + math.exp(-0.4 * (a - b)))
+            label = f"symbolic_{mode}"
         elif family == "mixed":
-            pick = random.choice(["housing", "subtraction", "multiplication", "symbolic"])
-            return self.generate_sample(pick)
+            sub = random.choice(["housing", "subtraction", "multiplication", "symbolic"])
+            return self._make_sample(sub)
         else:
-            a, b = random.uniform(2, 10), random.uniform(2, 10)
-            c, label = a + b, "default_add"
-        return {"a": float(a), "b": float(b), "c": float(c), "label": label}
+            c = a + b
+            label = "default_add"
+
+        noise = random.uniform(-self.config.sample_noise, self.config.sample_noise)
+        return {"a": float(a), "b": float(b), "c": float(c + noise), "label": label}
 
-    def _apply_sample_to_cells(self, sample: Dict[str, Any], anchor_output: bool):
+    def generate_example(self, family: Optional[str] = None):
+        with self.lock:
+            return self._make_sample(family or self.config.dataset_family)
+
+    def _apply_sample(self, sample: Dict[str, Any], anchor_output: bool):
         self.current_sample = sample
         self.last_sample = sample
-
         for cell in self.cells:
             cell.a = float(sample["a"])
             cell.b = float(sample["b"])
             cell.target = float(sample["c"]) if sample.get("c") is not None else None
             cell.label = sample.get("label", "")
             cell.anchored = anchor_output
-
-            if anchor_output:
-                cell.c = float(sample["c"])
-            else:
-                cell.c = 0.0
-
+            cell.c = float(sample["c"]) if anchor_output and sample.get("c") is not None else 0.0
             cell.prediction = 0.0
             cell.error = 0.0
-            cell.energy = 0.0
+            cell.loss = 0.0
             cell.force = 0.0
 
     def load_sample(self, sample: Dict[str, Any], anchor_output: Optional[bool] = None):
-        with self.lock:
-            if anchor_output is None:
-                anchor_output = self.config.mode == "training"
-            self._apply_sample_to_cells(sample, anchor_output=anchor_output)
-            self.add_log(
-                f"Sample loaded: a={sample['a']:.3f}, b={sample['b']:.3f}, "
-                f"c={sample['c']:.3f}, label={sample.get('label', '')}"
-            )
+        if anchor_output is None:
+            anchor_output = self.config.mode == "training" and sample.get("c") is not None
+        self._apply_sample(sample, anchor_output)
+        self.add_log(
+            f"Loaded sample: a={sample['a']:.3f}, b={sample['b']:.3f}, "
+            f"c={sample.get('c', None)}, label={sample.get('label', '')}"
+        )
 
-    def _coefficient_snapshot(self, cell: CellState):
-        if self.config.coeff_mode == "single_k":
-            return {"ka": cell.k, "kb": cell.k, "kc": cell.k}
-        if self.config.coeff_mode == "per_edge_k":
-            return {"ka": cell.ka, "kb": cell.kb, "kc": 0.0}
-        return {"ka": cell.ka, "kb": cell.kb, "kc": cell.kc}
-
-    def _set_trainable_param(self, cell: CellState, name: str, value: float):
-        value = max(-20.0, min(20.0, value))
-        if name == "k":
-            cell.k = value
-        elif name == "ka":
-            cell.ka = value
-        elif name == "kb":
-            cell.kb = value
-        elif name == "kc":
-            cell.kc = value
-
-    def _get_trainable_params(self):
-        if self.config.coeff_mode == "single_k":
-            return ["k"]
-        if self.config.coeff_mode == "per_edge_k":
-            return ["ka", "kb"]
-        return ["ka", "kb", "kc"]
+    def _neighbors(self, idx: int):
+        if self.config.topology == "single_cell":
+            return []
+        if self.config.topology == "chain":
+            n = []
+            if idx - 1 >= 0:
+                n.append(idx - 1)
+            if idx + 1 < len(self.cells):
+                n.append(idx + 1)
+            return n
+        return [j for j in range(len(self.cells)) if j != idx]
 
-    def _predict_cell(self, cell: CellState) -> float:
-        coeffs = self._coefficient_snapshot(cell)
+    def _predict(self, cell: CellState) -> float:
         a, b = cell.a, cell.b
-        arch = self.config.architecture
-        ka, kb, kc = coeffs["ka"], coeffs["kb"], coeffs["kc"]
 
         if self.config.coeff_mode == "single_k":
             k = cell.k
-            if arch == "additive":
+            if self.config.architecture == "additive":
                 return k * (a + b)
-            if arch == "multiplicative":
+            if self.config.architecture == "multiplicative":
                 return k * (a * b)
-            if arch == "affine":
-                return (k * a) + (k * b) + k
-            if arch == "bilinear":
-                return k * (a + b + (a * b))
-            if arch == "gated":
+            if self.config.architecture == "affine":
+                return k * a + k * b + k
+            if self.config.architecture == "gated":
                 gate = 1.0 / (1.0 + math.exp(-k))
                 return gate * (a + b) + (1.0 - gate) * (a * b)
             return k * (a + b)
 
-        if arch == "additive":
+        ka, kb, kc = cell.ka, cell.kb, cell.kc
+        if self.config.architecture == "additive":
             return (ka * a) + (kb * b) + kc
-        if arch == "multiplicative":
+        if self.config.architecture == "multiplicative":
             return (ka * a) * (kb * b) + kc
-        if arch == "affine":
+        if self.config.architecture == "affine":
             return (ka * a) + (kb * b) + kc
-        if arch == "bilinear":
-            return (ka * a) + (kb * b) + (kc * a * b)
-        if arch == "gated":
+        if self.config.architecture == "gated":
             gate = 1.0 / (1.0 + math.exp(-kc))
             return gate * ((ka * a) + (kb * b)) + (1.0 - gate) * (a * b)
         return (ka * a) + (kb * b) + kc
 
-    def _neighbors(self, idx: int):
-        if self.config.topology == "single_cell":
-            return []
-        if self.config.topology == "chain":
-            n = []
-            if idx - 1 >= 0:
-                n.append(idx - 1)
-            if idx + 1 < len(self.cells):
-                n.append(idx + 1)
-            return n
-        if self.config.topology == "mesh":
-            return [j for j in range(len(self.cells)) if j != idx]
-        return []
-
-    def _cell_loss(self, idx: int, preds: List[float]) -> float:
-        cell = self.cells[idx]
-        pred = preds[idx]
-        loss = 0.0
-        if cell.target is not None:
-            loss += (pred - cell.target) ** 2
-
-        neighbors = self._neighbors(idx)
-        if neighbors:
-            neighbor_mean = sum(preds[j] for j in neighbors) / len(neighbors)
-            loss += self.config.coupling * ((pred - neighbor_mean) ** 2)
-
-        return loss
-
-    def _numeric_gradient(self, idx: int, param_name: str, eps: float = 1e-4) -> float:
-        cell = self.cells[idx]
-        old = getattr(cell, param_name)
-
-        def local_loss() -> float:
-            pred = self._predict_cell(cell)
-            loss = 0.0
-            if cell.target is not None:
-                loss += (pred - cell.target) ** 2
-            neighbors = self._neighbors(idx)
-            if neighbors:
-                neighbor_preds = [self._predict_cell(self.cells[j]) for j in neighbors]
-                neighbor_mean = sum(neighbor_preds) / len(neighbor_preds)
-                loss += self.config.coupling * ((pred - neighbor_mean) ** 2)
-            return loss
-
-        self._set_trainable_param(cell, param_name, old + eps)
-        plus = local_loss()
-
-        self._set_trainable_param(cell, param_name, old - eps)
-        minus = local_loss()
-
-        self._set_trainable_param(cell, param_name, old)
-        return (plus - minus) / (2.0 * eps)
-
-    def _mean(self, xs: List[float]) -> float:
-        return sum(xs) / max(1, len(xs))
-
-    def _load_next_sample_from_batch(self):
+    def _train_cell(self, cell: CellState):
+        err = cell.prediction - (cell.target if cell.target is not None else 0.0)
+
+        if self.config.coeff_mode == "single_k":
+            if self.config.architecture == "additive":
+                grad = 2.0 * err * (cell.a + cell.b)
+            elif self.config.architecture == "multiplicative":
+                grad = 2.0 * err * (cell.a * cell.b)
+            elif self.config.architecture == "affine":
+                grad = 2.0 * err * (cell.a + cell.b + 1.0)
+            else:
+                gate = 1.0 / (1.0 + math.exp(-cell.k))
+                grad = 2.0 * err * ((cell.a + cell.b) - (cell.a * cell.b)) * gate * (1.0 - gate)
+
+            cell.k -= self.config.learning_rate * grad
+            cell.k = max(-20.0, min(20.0, cell.k))
+            return
+
+        if self.config.architecture == "additive":
+            da = 2.0 * err * cell.a
+            db = 2.0 * err * cell.b
+            dc = 2.0 * err * 1.0
+            cell.ka -= self.config.learning_rate * da
+            cell.kb -= self.config.learning_rate * db
+            cell.kc -= self.config.learning_rate * dc
+        elif self.config.architecture == "multiplicative":
+            da = 2.0 * err * (cell.kb * cell.b * cell.a)
+            db = 2.0 * err * (cell.ka * cell.a * cell.b)
+            dc = 2.0 * err * 1.0
+            cell.ka -= self.config.learning_rate * da
+            cell.kb -= self.config.learning_rate * db
+            cell.kc -= self.config.learning_rate * dc
+        elif self.config.architecture == "gated":
+            gate = 1.0 / (1.0 + math.exp(-cell.kc))
+            mix = (cell.ka * cell.a) + (cell.kb * cell.b)
+            da = 2.0 * err * gate * cell.a
+            db = 2.0 * err * gate * cell.b
+            dc = 2.0 * err * gate * (1.0 - gate) * (mix - (cell.a * cell.b))
+            cell.ka -= self.config.learning_rate * da
+            cell.kb -= self.config.learning_rate * db
+            cell.kc -= self.config.learning_rate * dc
+        else:
+            da = 2.0 * err * cell.a
+            db = 2.0 * err * cell.b
+            dc = 2.0 * err
+            cell.ka -= self.config.learning_rate * da
+            cell.kb -= self.config.learning_rate * db
+            cell.kc -= self.config.learning_rate * dc
+
+        cell.ka = max(-20.0, min(20.0, cell.ka))
+        cell.kb = max(-20.0, min(20.0, cell.kb))
+        cell.kc = max(-20.0, min(20.0, cell.kc))
+
+    def _pull_next_batch(self):
         if self.batch_queue:
-            sample = self.batch_queue.popleft()
-            self._apply_sample_to_cells(sample, anchor_output=(self.config.mode == "training"))
-            self.add_log(f"Next batch sample: {sample.get('label', '')}")
+            nxt = self.batch_queue.popleft()
+            self.load_sample(nxt, anchor_output=(self.config.mode == "training" and nxt.get("c") is not None))
             return True
         return False
 
-    def physics_step(self):
-        with self.lock:
-            if not self.running or not self.cells:
-                return False
+    def step(self):
+        if not self.running or not self.cells:
+            return False
 
-            preds = []
-            for cell in self.cells:
-                pred = self._predict_cell(cell)
-                cell.prediction = pred
-                preds.append(pred)
+        preds = []
+        for cell in self.cells:
+            cell.prediction = self._predict(cell)
+            preds.append(cell.prediction)
 
-            global_pred = self._mean(preds)
-            target_available = self.current_sample is not None and self.current_sample.get("c") is not None
-            target = self._mean([c.target for c in self.cells if c.target is not None]) if target_available else None
+        for idx, cell in enumerate(self.cells):
+            neigh = self._neighbors(idx)
+            neigh_mean = sum(preds[j] for j in neigh) / len(neigh) if neigh else cell.prediction
 
             if self.config.mode == "training":
-                total_loss = 0.0
-
-                for idx, cell in enumerate(self.cells):
-                    cell.target = target if target is not None else cell.target
-                    cell.error = (cell.prediction - cell.target) if cell.target is not None else 0.0
-                    cell.energy = cell.error ** 2
-
-                    for param_name in self._get_trainable_params():
-                        grad = self._numeric_gradient(idx, param_name)
-                        old = getattr(cell, param_name)
-                        new_val = old - (self.config.learning_rate * grad)
-                        new_val = (1.0 - self.config.damping) * new_val + (self.config.damping * old)
-                        self._set_trainable_param(cell, param_name, new_val)
-
-                    total_loss += self._cell_loss(idx, preds)
-
-                self.current_loss = total_loss / max(1, len(self.cells))
-                self.current_error = (global_pred - target) if target is not None else global_pred
-                self.loss_history.append(self.current_loss)
-                self.error_history.append(self.current_error)
-
-                if target_available and abs(self.current_error) < 0.05 and self.current_loss < 0.01:
-                    self.add_log("Converged on current sample.")
-                    if self._load_next_sample_from_batch():
-                        self.iteration += 1
-                        return True
-                    self.running = False
-                    self.add_log("Batch complete.")
-                    self.iteration += 1
-                    return False
+                if cell.target is None:
+                    continue
+
+                cell.error = cell.prediction - cell.target
+                cell.loss = cell.error ** 2 + self.config.coupling * ((cell.prediction - neigh_mean) ** 2)
+                self._train_cell(cell)
 
             else:
-                # Inference mode: output node(s) drift toward the predicted state.
-                drift_values = []
-                for idx, cell in enumerate(self.cells):
-                    neighbors = self._neighbors(idx)
-                    neighbor_mean = self._mean([preds[j] for j in neighbors]) if neighbors else pred
-
-                    drift = (pred - cell.c)
-                    drift += self.config.coupling * (neighbor_mean - cell.c)
-
-                    cell.force = drift
-                    cell.c += 0.15 * drift
-                    cell.error = pred - cell.c
-                    cell.energy = cell.error ** 2
-                    drift_values.append(abs(drift))
-
-                self.current_error = self._mean([cell.error for cell in self.cells])
-                self.current_loss = self._mean([cell.energy for cell in self.cells])
-                self.loss_history.append(self.current_loss)
-                self.error_history.append(self.current_error)
-
-                if self.current_sample and abs(self.current_error) < 0.05 and self.current_loss < 0.01:
-                    self.add_log("Inference settled.")
-                    if self._load_next_sample_from_batch():
-                        self.iteration += 1
-                        return True
-                    self.running = False
-                    self.add_log("Task complete.")
-                    self.iteration += 1
-                    return False
-
-                # If no target exists, stop when drift is tiny.
-                if not target_available and self._mean(drift_values) < 0.002:
-                    self.running = False
-                    self.add_log("Inference drift stabilized.")
-                    self.iteration += 1
-                    return False
-
-            self.iteration += 1
-            return True
+                if cell.target is None:
+                    cell.force = cell.prediction - cell.c
+                    cell.c += 0.18 * cell.force
+                    cell.error = cell.prediction - cell.c
+                    cell.loss = cell.error ** 2
+                else:
+                    cell.force = (cell.prediction - cell.c) + self.config.coupling * (neigh_mean - cell.c)
+                    cell.c += 0.16 * cell.force
+                    cell.error = cell.prediction - cell.c
+                    cell.loss = cell.error ** 2
+
+        self.current_error = sum(c.error for c in self.cells) / len(self.cells)
+        self.current_loss = sum(c.loss for c in self.cells) / len(self.cells)
+        self.loss_history.append(self.current_loss)
+        self.error_history.append(self.current_error)
+
+        self.iteration += 1
+
+        converged = abs(self.current_error) < 0.05 and self.current_loss < 0.01
+        if converged:
+            if self._pull_next_batch():
+                self.add_log("Switched to next batch item.")
+                return True
+            self.running = False
+            self.add_log("Converged.")
+            return False
+
+        return True
 
     def start_batch(self, count: int):
-        with self.lock:
-            self.batch_queue.clear()
-            for _ in range(count):
-                self.batch_queue.append(self.generate_sample())
-            first = self._load_next_sample_from_batch()
-            self.running = first
-            self.add_log(f"Batch started with {count} samples.")
-            return first
-
-    def set_custom_sample(self, a: float, b: float, c: Optional[float] = None):
-        with self.lock:
-            sample = {"a": float(a), "b": float(b), "c": float(c) if c is not None else None, "label": "custom"}
-            self._apply_sample_to_cells(sample, anchor_output=(self.config.mode == "training" and c is not None))
-            self.current_sample = sample
-            self.last_sample = sample
-            self.running = True
-            self.add_log(f"Custom sample loaded: a={a}, b={b}, c={c}")
-            return sample
+        self.batch_queue.clear()
+        for _ in range(count):
+            self.batch_queue.append(self.generate_example())
+        first = self.batch_queue.popleft()
+        self.load_sample(first, anchor_output=(self.config.mode == "training"))
+        self.running = True
+        self.add_log(f"Batch started: {count} samples.")
+        return first
+
+    def custom_input(self, a: float, b: float, c: Optional[float]):
+        sample = {"a": float(a), "b": float(b), "c": float(c) if c is not None else None, "label": "custom"}
+        self.load_sample(sample, anchor_output=(self.config.mode == "training" and c is not None))
+        self.running = True
+        self.add_log("Custom input loaded.")
+        return sample
 
     def halt(self):
-        with self.lock:
-            self.running = False
-            self.add_log("Engine halted.")
-
-    def snapshot(self) -> Dict[str, Any]:
-        with self.lock:
-            return {
-                "config": asdict(self.config),
-                "running": self.running,
-                "iteration": self.iteration,
-                "current_error": self.current_error,
-                "current_loss": self.current_loss,
-                "cells": [c.to_dict() for c in self.cells],
-                "logs": self.logs,
-                "last_sample": self.last_sample,
-                "current_sample": self.current_sample,
-                "batch_remaining": len(self.batch_queue),
-                "loss_history": list(self.loss_history),
-                "error_history": list(self.error_history),
-            }
+        self.running = False
+        self.add_log("Halted.")
+
+    def snapshot(self):
+        return {
+            "config": asdict(self.config),
+            "running": self.running,
+            "iteration": self.iteration,
+            "current_error": self.current_error,
+            "current_loss": self.current_loss,
+            "last_action": self.last_action,
+            "cells": [c.to_dict() for c in self.cells],
+            "logs": self.logs,
+            "current_sample": self.current_sample,
+            "last_sample": self.last_sample,
+            "batch_remaining": len(self.batch_queue),
+            "loss_history": list(self.loss_history),
+            "error_history": list(self.error_history),
+        }
 
 
 engine = SimEngine()
 
 
-def run_loop():
+def loop():
     while True:
-        if engine.running:
-            engine.physics_step()
+        try:
+            if engine.running:
+                engine.step()
+        except Exception as e:
+            with engine.lock:
+                engine.running = False
+                engine.add_log(f"Engine error: {e}")
         time.sleep(0.04)
 
 
-threading.Thread(target=run_loop, daemon=True).start()
+threading.Thread(target=loop, daemon=True).start()
 
 
-@app.get("/", response_class=HTMLResponse)
-async def get_ui():
-    return FileResponse("index.html")
+@app.get("/")
+async def root():
+    return FileResponse(INDEX_PATH)
 
 
 @app.get("/state")
-async def get_state():
-    return engine.snapshot()
+async def state():
+    return JSONResponse(engine.snapshot())
 
 
 @app.post("/config")
-async def config(data: dict):
-    engine.configure(data)
-    return {"ok": True}
+async def config(req: Request):
+    payload = await req.json()
+    engine.configure(payload)
+    return {"ok": True, "message": "Config applied."}
 
 
 @app.post("/example")
-async def example(data: dict):
-    family = data.get("dataset_family", engine.config.dataset_family)
-    sample = engine.generate_sample(family)
-    return sample
+async def example(req: Request):
+    payload = await req.json()
+    family = payload.get("dataset_family", engine.config.dataset_family)
+    return engine.generate_example(family)
 
 
 @app.post("/generate_batch")
-async def generate_batch(data: dict):
-    count = int(data.get("count", engine.config.batch_size))
+async def generate_batch(req: Request):
+    payload = await req.json()
+    count = int(payload.get("count", engine.config.batch_size))
     engine.start_batch(count)
-    return {"ok": True, "count": count}
+    return {"ok": True, "message": f"Batch started with {count} samples."}
 
 
 @app.post("/test_custom")
-async def test_custom(data: dict):
-    a = float(data["a"])
-    b = float(data["b"])
-    c = data.get("c", None)
-    c_val = float(c) if c not in [None, "", "null"] else None
-    engine.set_custom_sample(a, b, c_val)
-    return {"ok": True}
+async def test_custom(req: Request):
+    payload = await req.json()
+    a = float(payload["a"])
+    b = float(payload["b"])
+    c_val = payload.get("c", None)
+    c = None if c_val in [None, "", "null"] else float(c_val)
+    engine.custom_input(a, b, c)
+    return {"ok": True, "message": "Custom input loaded."}
 
 
 @app.post("/halt")
 async def halt():
     engine.halt()
-    return {"ok": True}
+    return {"ok": True, "message": "Engine halted."}
 
 
 if __name__ == "__main__":