import os import time import json import asyncio import threading import numpy as np from typing import Dict, Any, Optional, List, Callable from src.connectome.types import ConnectomeGraph, GraphMode, ProvenanceStatus from src.connectome.loader import get_or_create_circuit from src.brain.runtime import BrainRuntime from src.memory.persistence import PersistentMemoryManager from src.trainer.curriculum import CurriculumTrainer from src.evolution.scheduler import EvolutionScheduler from src.dream.engine import DreamEngine class SimulationEngine: """ Single Authoritative Simulation Engine. Owns the central simulation loop, manages thread-safe state access, dispatches telemetry to connected WebSocket clients, and coordinates subsystems (brain, memory, learning, evolution, dreaming). """ def __init__( self, circuit_size: int = 512, graph_mode: GraphMode = GraphMode.REAL, use_gpu: bool = True, seed: int = 42, db_path: str = "diagnostics/flybrain_live_memory.db" ): self.lock = threading.RLock() self.seed = seed self.graph_mode = graph_mode self.circuit_size = circuit_size self.use_gpu = use_gpu self.db_path = db_path # Initialize subsystems self.circuit = get_or_create_circuit(circuit_size, mode=graph_mode, seed=seed) self.memory = PersistentMemoryManager(db_path=db_path) self.brain = BrainRuntime(self.circuit, use_gpu=use_gpu, seed=seed) self.trainer = CurriculumTrainer(self.brain, self.memory) self.evolution = EvolutionScheduler(self.circuit, history_file="diagnostics/evolution_history.json") self.dream_engine = DreamEngine(self.brain, self.memory) # Concurrency & Execution Control self.is_running = False self.target_hz = 10.0 self._stop_event = threading.Event() self._worker_thread: Optional[threading.Thread] = None # Sensory Stimulus Queue self.pending_sensory: Dict[str, np.ndarray] = {} self.pending_reward: float = 0.0 # Telemetry broadcast queues (for async WebSockets) self.telemetry_listeners: List[asyncio.Queue] = [] self._loop: Optional[asyncio.AbstractEventLoop] = None # Performance Telemetry self.last_step_time_ms = 0.0 self.step_history: List[float] = [] # Stream mode (V5 phase_10): 24/7 supervision surface. self.boot_time = time.time() self.stream_mode = "STOPPED" # STOPPED | RUNNING | PAUSED self.heartbeat = {"ts": time.time(), "step": 0, "note": "boot"} def set_event_loop(self, loop: asyncio.AbstractEventLoop): self._loop = loop def start(self): with self.lock: if not self.is_running: self.is_running = True self.stream_mode = "RUNNING" self._stop_event.clear() self._worker_thread = threading.Thread(target=self._run_loop, daemon=True, name="FlyBrainSimWorker") self._worker_thread.start() def pause(self): with self.lock: self.is_running = False if self.stream_mode == "RUNNING": self.stream_mode = "PAUSED" self._stop_event.set() if self._worker_thread and self._worker_thread.is_alive(): self._worker_thread.join(timeout=1.0) self._worker_thread = None def resume(self) -> Dict[str, Any]: """Resume a paused stream (no-op when already running).""" with self.lock: was = self.stream_mode if was == "PAUSED": self.start() return {"status": "RESUMED" if was == "PAUSED" else "NOOP", "previous": was, "stream_mode": self.stream_mode} def stop(self) -> Dict[str, Any]: """Stop the stream loop; simulation state is preserved.""" self.pause() with self.lock: self.stream_mode = "STOPPED" return {"status": "STOPPED", "step": self.brain.state.step_count} def safe_shutdown(self) -> Dict[str, Any]: """Stop loop + close GPU; state stays in memory for backup/snapshot.""" res = self.stop() try: self.brain.cleanup() except Exception: pass return {**res, "status": "SAFE_SHUTDOWN"} def restart_runtime(self) -> Dict[str, Any]: """Destructive fallback: rebuild BrainRuntime on the same circuit (fresh neural state). Prefer restore-from-backup; this is the watchdog's third recovery rung, always logged by the caller.""" self.pause() try: self.brain.cleanup() except Exception: pass self.brain = BrainRuntime(self.circuit, use_gpu=self.use_gpu, seed=self.seed) self.last_step_time_ms = 0.0 self.step_history = [] return {"status": "RUNTIME_RESTARTED", "step": 0} def stream_status(self) -> Dict[str, Any]: """24/7 stream status: uptime, sim time, population vitals, backend.""" with self.lock: now = time.time() hb = dict(self.heartbeat) hb.update({"ts": now, "step": self.brain.state.step_count, "stream_mode": self.stream_mode, "is_running": self.is_running}) self.heartbeat = dict(hb) lat = list(self.step_history) return { "stream_mode": self.stream_mode, "is_running": self.is_running, "uptime_sec": round(now - self.boot_time, 1), "simulation_step": self.brain.state.step_count, "simulation_time_note": "1 step per loop iteration at target_hz", "target_hz": self.target_hz, "active_spikes": int(np.sum(self.brain.state.spikes > 0.5)), "total_spikes": self.brain.state.total_spikes, "mean_latency_ms": round(float(np.mean(lat)), 3) if lat else 0.0, "p95_latency_ms": round(float(np.percentile(lat, 95)), 3) if lat else 0.0, "backend": "vulkan_gpu" if (self.brain.gpu_engine and self.brain.use_gpu) else "cpu_reference", "heartbeat": hb, } def _run_loop(self): """Authoritative single simulation thread.""" while not self._stop_event.is_set() and self.is_running: delay = 1.0 / max(1.0, self.target_hz) t0 = time.perf_counter() # Execute step under lock with self.lock: sensory = dict(self.pending_sensory) if self.pending_sensory else None self.pending_sensory.clear() reward = self.pending_reward self.pending_reward = 0.0 res = self.brain.step(sensory_inputs=sensory, reward=reward) t1 = time.perf_counter() self.last_step_time_ms = round((t1 - t0) * 1000.0, 3) self.step_history.append(self.last_step_time_ms) if len(self.step_history) > 100: self.step_history.pop(0) # Broadcast telemetry to async listeners self._broadcast_telemetry(res) elapsed = time.perf_counter() - t0 sleep_time = max(0.001, delay - elapsed) time.sleep(sleep_time) def step_single( self, n_steps: int = 1, sensory_inputs: Optional[Dict[str, np.ndarray]] = None, reward: float = 0.0 ) -> Dict[str, Any]: """Manually steps the simulation by N steps under lock.""" with self.lock: res = {} for _ in range(n_steps): t0 = time.perf_counter() res = self.brain.step(sensory_inputs=sensory_inputs, reward=reward) t1 = time.perf_counter() self.last_step_time_ms = round((t1 - t0) * 1000.0, 3) self.step_history.append(self.last_step_time_ms) if len(self.step_history) > 100: self.step_history.pop(0) self._broadcast_telemetry(res) return res def set_sensory_stimulus(self, population_name: str, values: np.ndarray): with self.lock: self.pending_sensory[population_name] = values def add_reward(self, reward: float): with self.lock: self.pending_reward += reward def register_telemetry_queue(self, queue: asyncio.Queue): with self.lock: self.telemetry_listeners.append(queue) def unregister_telemetry_queue(self, queue: asyncio.Queue): with self.lock: if queue in self.telemetry_listeners: self.telemetry_listeners.remove(queue) def _broadcast_telemetry(self, step_result: Dict[str, Any]): """Dispatches lightweight telemetry to registered WebSocket queues.""" payload = self.get_telemetry_payload() if self._loop and self.telemetry_listeners: for q in list(self.telemetry_listeners): try: self._loop.call_soon_threadsafe(q.put_nowait, payload) except Exception: pass def get_telemetry_payload(self) -> Dict[str, Any]: """Returns compact real-time telemetry dictionary without full connectome flooding.""" with self.lock: brain = self.brain active_spikes = int(np.sum(brain.state.spikes > 0.5)) top_active = [ {"idx": int(i), "act": round(float(brain.state.activations[i]), 2)} for i in np.argsort(brain.state.activations)[-10:] ] gpu_diag = brain.gpu_engine.get_diagnostics() if brain.gpu_engine else {} return { "step": brain.state.step_count, "spikes": active_spikes, "total_spikes": brain.state.total_spikes, "mean_activation": round(float(np.mean(brain.state.activations)), 4), "mean_potential": round(float(np.mean(brain.state.membrane_potentials)), 4), "prediction_error": round(brain.state.prediction_error, 4), "predicted_reward": round(brain.state.predicted_reward, 4), "current_reward": round(brain.state.current_reward, 4), "drives": { "energy": round(brain.state.drives.energy, 3), "curiosity": round(brain.state.drives.curiosity, 3), "social": round(brain.state.drives.social, 3), "integrity": round(brain.state.drives.integrity, 3) }, "tool_associations": dict(brain.state.tool_associations), "active_neurons": top_active, "step_latency_ms": self.last_step_time_ms, "backend": "vulkan_gpu" if (brain.gpu_engine and brain.use_gpu) else "cpu_reference", "device_name": brain.gpu_engine.device_name if brain.gpu_engine else "CPU Reference", "is_running": self.is_running, "graph_mode": self.circuit.mode.value, "graph_hash": self.circuit.graph_hash[:16] } def get_full_state(self) -> Dict[str, Any]: """Returns complete state for diagnostics and dashboard inspection.""" with self.lock: return { "simulation": self.get_telemetry_payload(), "connectome": { "num_neurons": self.circuit.num_neurons, "num_synapses": self.circuit.num_synapses, "mode": self.circuit.mode.value, "provenance_status": self.circuit.provenance_status.value, "graph_hash": self.circuit.graph_hash, "populations": self.circuit.populations.to_dict() if self.circuit.populations else {} }, "gpu_diagnostics": self.brain.gpu_engine.get_diagnostics() if self.brain.gpu_engine else { "status": "UNAVAILABLE", "device_name": "CPU Reference Mode" } } def get_connectome_3d_view(self, max_nodes: int = 512, max_edges: int = 384) -> Dict[str, Any]: """ Returns normalized 3D coordinates, biological soma metadata, and active synaptic edges for the interactive 3D WebGL viewer. """ with self.lock: graph = self.circuit N = min(max_nodes, graph.num_neurons) coords = graph.coordinates[:N] c_min = coords.min(axis=0) c_max = coords.max(axis=0) norm_coords = (coords - c_min) / (c_max - c_min + 1e-5) * 2.0 - 1.0 # Normalized to [-1, 1] nodes = [] for i in range(N): nodes.append({ "id": int(graph.neuron_ids[i]), "idx": i, "pos": [round(float(norm_coords[i, 0]), 3), round(float(norm_coords[i, 1]), 3), round(float(norm_coords[i, 2]), 3)], "side": graph.sides[i], "tbars": int(graph.tbars[i]), "act": round(float(self.brain.state.activations[i]), 3), "pot": round(float(self.brain.state.membrane_potentials[i]), 3), "spk": int(self.brain.state.spikes[i] > 0.5) }) edges = [] edge_count = 0 # CSR v3: row i stores INCOMING sources; displayed edge = source -> i. for i in range(min(128, N)): start = graph.row_offsets[i] end = min(start + 4, graph.row_offsets[i + 1]) for k in range(start, end): src = int(graph.col_indices[k]) if src < N: edges.append({ "src": src, "tgt": i, "w": round(float(graph.weights[k]), 2) }) edge_count += 1 if edge_count >= max_edges: break if edge_count >= max_edges: break return { "num_neurons": graph.num_neurons, "num_synapses": graph.num_synapses, "mode": graph.mode.value, "provenance_status": graph.provenance_status.value, "graph_hash": graph.graph_hash, "nodes": nodes, "edges": edges } def close(self): self.pause() self.brain.cleanup() def cleanup(self): self.close() def get_latest_telemetry(self) -> Dict[str, Any]: return self.get_telemetry_payload() def get_current_state(self) -> Dict[str, Any]: return self.get_full_state()