From eba2592aa0f97e51e5c2d6e7d47f653f0762667d Mon Sep 17 00:00:00 2001
From: Capek System Safety & Robotic Solutions
 <96583804+MelkorBalrog@users.noreply.github.com>
Date: Thu, 7 Aug 2025 10:07:38 -0400
Subject: [PATCH] Add tests for class names and GET_NUM_CLASSES?

---
 synapse/hardware/cpu.py        |  7 ++++---
 synapse/models/redundant_ip.py | 19 +++++++++++++++++++
 tests/test_cpu_class_labels.py | 27 +++++++++++++++++++++++++++
 3 files changed, 50 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_cpu_class_labels.py

diff --git a/synapse/hardware/cpu.py b/synapse/hardware/cpu.py
index 73d3653..ec3a595 100644
--- a/synapse/hardware/cpu.py
+++ b/synapse/hardware/cpu.py
@@ -69,8 +69,9 @@ def step(self, program):
             self.running = False
             if self.neural_ip.last_result is not None:
                 result = self.get_reg("$t9")
-                label_map = {0: "A", 1: "B", 2: "Unknown"}
-                print(f"Final classification: {label_map.get(result, result)}")
+                class_names = getattr(self.neural_ip, "class_names", [])
+                label = class_names[result] if 0 <= result < len(class_names) else result
+                print(f"Final classification: {label}")
         elif instr == "ADDI":
             rd = parts[1].rstrip(",")
             rs = parts[2].rstrip(",")
@@ -99,6 +100,6 @@ def step(self, program):
         elif instr == "OP_NEUR":
             subcmd = " ".join(parts[1:])
             self.neural_ip.run_instruction(subcmd, memory=self.memory)
-            if subcmd.upper().startswith("INFER_ANN") and self.neural_ip.last_result is not None:
+            if any(subcmd.upper().startswith(cmd) for cmd in ("INFER_ANN", "GET_ARGMAX", "GET_NUM_CLASSES")) and self.neural_ip.last_result is not None:
                 self.set_reg("$t9", int(self.neural_ip.last_result))
 
diff --git a/synapse/models/redundant_ip.py b/synapse/models/redundant_ip.py
index 66fbd18..8d421d8 100644
--- a/synapse/models/redundant_ip.py
+++ b/synapse/models/redundant_ip.py
@@ -55,6 +55,12 @@ def __init__(self, train_data_dir: str | None = None) -> None:
         # Metrics and figures generated during training keyed by ANN ID
         self.metrics_by_ann: Dict[int, Dict[str, float]] = {}
         self.figures_by_ann: Dict[int, List] = {}
+        # Class names inferred from the training directory for label mapping
+        self.class_names: List[str] = []
+        if train_data_dir:
+            root = Path(train_data_dir)
+            if root.is_dir():
+                self.class_names = sorted([d.name for d in root.iterdir() if d.is_dir()])
 
     # ------------------------------------------------------------------
     # Assembly interface
@@ -84,6 +90,15 @@ def run_instruction(self, subcmd: str, memory=None) -> None:
                     ann.load(f"{prefix}_{ann_id}.pt")
                 except FileNotFoundError:
                     pass
+        elif op == "GET_ARGMAX":
+            # Argmax already computed during INFER_ANN; expose via last_result
+            pass
+        elif op == "GET_NUM_CLASSES":
+            if not self.class_names and self.train_data_dir:
+                root = Path(self.train_data_dir)
+                if root.is_dir():
+                    self.class_names = sorted([d.name for d in root.iterdir() if d.is_dir()])
+            self.last_result = len(self.class_names) if self.class_names else hp.num_classes
         elif op == "SAVE_PROJECT":
             json_path = tokens[1] if len(tokens) > 1 else "project.json"
             prefix = tokens[2] if len(tokens) > 2 else "weights"
@@ -231,6 +246,10 @@ def predict_majority(self, X: np.ndarray):
     # Internal helpers
     # ------------------------------------------------------------------
     def _load_dataset(self):
+        if not self.class_names and self.train_data_dir:
+            root = Path(self.train_data_dir)
+            if root.is_dir():
+                self.class_names = sorted([d.name for d in root.iterdir() if d.is_dir()])
         if self._cached_dataset is None:
             if not self.train_data_dir:
                 print("No training data directory specified; aborting training.")
diff --git a/tests/test_cpu_class_labels.py b/tests/test_cpu_class_labels.py
new file mode 100644
index 0000000..4242041
--- /dev/null
+++ b/tests/test_cpu_class_labels.py
@@ -0,0 +1,27 @@
+import os
+import sys
+
+import pytest
+
+sys.path.append(os.getcwd())
+from synapse.models.redundant_ip import RedundantNeuralIP
+from synapse.hardware.cpu import CPU
+
+
+def test_get_num_classes(tmp_path):
+    (tmp_path / "class0").mkdir()
+    (tmp_path / "class1").mkdir()
+    ip = RedundantNeuralIP(train_data_dir=str(tmp_path))
+    ip.run_instruction("GET_NUM_CLASSES")
+    assert ip.last_result == 2
+
+
+def test_cpu_halt_uses_class_names(capsys):
+    ip = RedundantNeuralIP()
+    ip.class_names = ["class0", "class1"]
+    ip.last_result = 1
+    cpu = CPU("CPU1", None, ip, None, None)
+    cpu.set_reg("$t9", 1)
+    cpu.step(["HALT"])
+    captured = capsys.readouterr().out
+    assert "Final classification: class1" in captured