wcm_storage.py

"""WCM Storage Mixin — bistable grid storage, content-addressable grep, sprite blit.

Provides StorageMixin, the persistence layer composed into WaveKernel.
Covers bistable-latch storage (bits encoded as stable V_HIGH / V_LOW fixed
points on the FHN grid), zone allocation with overdamped guard bands,
content-addressable grep via analytic Hamming-distance probes, and sprite
blit for stamping stored patterns onto the voltage surface.

Architectural role (see docs/DEV_MODE.md §2.1):
  This is a chipset storage primitive — it provides the *how* of persistence
  (voltage latches, zone layout, probe physics).  All decisions about *what*
  to store or *when* to persist are made by WaveIL policy, not here.
"""

import numpy as np
from typing import List, Optional, Tuple

from wcm_constants import (
	A_FHN,
	A_GUARD,
	A_STORE,
	B_FHN,
	B_GUARD,
	B_STORE,
	BISTABLE_THRESHOLD,
	CHIP_REGION_DESKTOP_ROWS,
	DIFFUSE,
	EPS,
	EPS_GUARD,
	EPS_STORE,
	GLYPH_ATLAS_CELL_H,
	GLYPH_ATLAS_ROW0,
	GLYPH_ATLAS_ROWS_OF_GLYPHS,
	STORAGE_PROBE_AMP,
	V_BOUNDARY_HIGH,
	V_BOUNDARY_LOW,
	V_HIGH,
	V_LOW,
	V_REST,
	W_HIGH,
	W_LOW,
	W_REST,
	AttractorObject,
)


class StorageMixin:
	"""Bistable storage primitives mixed into WaveKernel.

	Responsibilities:
	  - Zone lifecycle: allocate / deallocate rectangular storage regions
	    with guard-band isolation on the FHN grid.
	  - Pattern I/O: save and load domain voltage snapshots as bistable
	    bit-patterns (wave-native primary path) or dict payloads (legacy).
	  - Content-addressable grep: Hamming-distance probe across stored
	    slots without reading bits back to Python.
	  - Sprite blit: stamp a stored pattern onto the voltage grid or an
	    off-grid float buffer for display compositing.

	Assumed host state (provided by WaveKernel / other mixins):
	  self.v, self.w          — voltage and recovery grids (float32)
	  self.rows, self.cols    — grid dimensions
	  self.a_map, self.b_map, self.eps_map, self.d — FHN parameter arrays
	  self.param_zones        — dict of named zones (managed here)
	  self.storage_freeze_mask — bool grid; True = excluded from sparse frontier
	  self.store              — AttractorStore for metadata / legacy payloads
	  self.wave_native_storage — bool; True = grid-latch primary path
	  self.domains            — domain registry (from SubstrateMixin)
	  self.domain_to_pid      — domain→pid lookup
	  self.tick_count         — current simulation tick
	"""

	# --- Pattern object API (save / load / GC) --------------------------------

	def _ensure_domain_storage_zone(self, domain_name: str) -> Optional[str]:
		"""Auto-allocate a grid storage zone for a domain if wave_native_storage is on.

		Places the zone below the desktop/glyph-atlas region, stacking
		vertically beneath any previously allocated zones.  Returns the
		zone name on success, None if the domain is missing or the grid
		has no room.
		"""
		zone_name = f"domain_store_{domain_name}"
		if zone_name in self.param_zones:
			return zone_name
		if domain_name not in self.domains:
			return None
		d = self.domains[domain_name]
		n_cells = int(np.count_nonzero(d.mask))
		if n_cells <= 0:
			return None
		# Place zone in non-desktop region below desktop rows
		side = int(np.ceil(np.sqrt(n_cells)))
		zone_rows = side
		zone_cols = int(np.ceil(n_cells / side))
		# Find a free row band below glyph atlas
		atlas_end = int(GLYPH_ATLAS_ROW0) + int(GLYPH_ATLAS_CELL_H) * int(GLYPH_ATLAS_ROWS_OF_GLYPHS) + 2
		base_r = max(int(CHIP_REGION_DESKTOP_ROWS) + 4, atlas_end)
		for zn, zinfo in self.param_zones.items():
			if zinfo.get("kind") == "storage":
				base_r = max(base_r, zinfo["r1"] + zinfo.get("guard_width", 2) + 2)
		if base_r + zone_rows + 4 > self.rows or zone_cols + 4 > self.cols:
			return None
		ok = self.allocate_storage_zone(zone_name, base_r, 0, zone_rows, zone_cols, guard_width=2)
		return zone_name if ok else None

	def save_pattern_object(
		self,
		key: str,
		domain_name: str,
		metadata: dict = None,
		rehydrate_policy: str = "blend",
		integrity_score: float = 1.0,
	) -> bool:
		self._bump_syscall("save_pattern_object", pid=self.domain_to_pid.get(domain_name))
		if domain_name not in self.domains:
			return False
		d = self.domains[domain_name]
		v_domain = self.v[d.mask]

		# Wave-native primary path: voltage → bistable latch.
		# This is the target storage model — domain voltages are thresholded
		# into bits and written as V_HIGH / V_LOW latches in a dedicated
		# storage zone.  The dict-based fallback below is legacy and will
		# be removed once all callers migrate to wave-native.
		if self.wave_native_storage:
			zone_name = self._ensure_domain_storage_zone(domain_name)
			if zone_name is not None:
				bits = (v_domain > np.float32(BISTABLE_THRESHOLD)).astype(np.int32)
				zone = self.param_zones[zone_name]
				pr, pc = zone["patch_rows"], zone["patch_cols"]
				slot_size = pr * pc
				if bits.size < slot_size:
					bits = np.pad(bits, (0, slot_size - bits.size), constant_values=0)
				else:
					bits = bits[:slot_size]
				grid_ok = self.save_to_grid(key, bits, zone_name)
				# Metadata sidecar only — payload not authoritative
				self.store.save_object(AttractorObject(
					id=key,
					metadata=dict(metadata or {"domain": domain_name, "storage": "grid"}),
					payload=np.zeros(0, dtype=np.float32),
					mask=np.array(d.mask, copy=True),
					rehydrate_policy=rehydrate_policy,
					integrity_score=float(integrity_score),
					created_tick=int(self.tick_count),
				))
				return grid_ok

		# Fallback: dict-based storage (legacy path)
		obj = AttractorObject(
			id=key,
			metadata=dict(metadata or {"domain": domain_name}),
			payload=np.array(v_domain, copy=True),
			mask=np.array(d.mask, copy=True),
			rehydrate_policy=rehydrate_policy,
			integrity_score=float(integrity_score),
			created_tick=int(self.tick_count),
		)
		return self.store.save_object(obj)

	def _pattern_similarity(self, a: np.ndarray, b: np.ndarray) -> float:
		"""Return 0-1 similarity between two voltage arrays (1 = identical).

		Uses mean absolute difference normalised to the clipped voltage
		range [-5, 5], so max possible difference of 10 maps to 0.0.
		"""
		if a.size == 0 or b.size == 0 or a.shape != b.shape:
			return 0.0
		diff = np.mean(np.abs(a.astype(np.float32) - b.astype(np.float32)))
		# v is clipped to [-5,5], so normalized max difference is 10.
		return float(max(0.0, min(1.0, 1.0 - (diff / 10.0))))

	def load_pattern_object(self, key: str, domain_name: str, gain: float = 0.7):
		"""Load pattern into a domain — grid-native primary, dict fallback."""
		self._bump_syscall("load_pattern_object", pid=self.domain_to_pid.get(domain_name))
		if domain_name not in self.domains:
			return False, 0.0
		d = self.domains[domain_name]
		n_domain = int(np.count_nonzero(d.mask))

		# Wave-native primary: read latches, decode, blend into domain
		if self.wave_native_storage:
			zone_name = f"domain_store_{domain_name}"
			if zone_name in self.param_zones:
				bits = self.load_from_grid(key, zone_name)
				if bits is not None:
					bits = bits[:n_domain]
					v_restored = np.where(bits, np.float32(V_HIGH), np.float32(V_LOW))
					self.v[d.mask] = self.v[d.mask] * np.float32(1.0 - gain) + v_restored * np.float32(gain)
					after = np.array(self.v[d.mask], copy=True)
					return True, self._pattern_similarity(after, v_restored)

		# Dict fallback
		obj = self.store.load_object(key)
		if obj is None:
			return False, 0.0
		self.v[d.mask] = self.v[d.mask] * np.float32(1.0 - gain) + obj.payload * np.float32(gain)
		after = np.array(self.v[d.mask], copy=True)
		return True, self._pattern_similarity(after, obj.payload)

	def gc_patterns(self, min_integrity: float = 0.0, keep_latest: int = 256) -> int:
		self._bump_syscall("gc_patterns")
		return self.store.gc(min_integrity=min_integrity, keep_latest=keep_latest)

	# --- Wave-native bistable storage (zone lifecycle) -----------------------

	def allocate_storage_zone(self, name: str, r0: int, c0: int, zone_rows: int, zone_cols: int, guard_width: int = 2, patch_rows: int = 0, patch_cols: int = 0) -> bool:
		"""Allocate a rectangular storage zone with bistable parameters and guard band.
		
		If patch_rows/patch_cols > 0, the zone is subdivided into a grid of
		fixed-size slots, each independently addressable by key.
		"""
		if name in self.param_zones:
			return False
		r1 = r0 + zone_rows
		c1 = c0 + zone_cols
		if r0 < 0 or c0 < 0 or r1 > self.rows or c1 > self.cols:
			return False
		gr0 = max(0, r0 - guard_width)
		gr1 = min(self.rows, r1 + guard_width)
		gc0 = max(0, c0 - guard_width)
		gc1 = min(self.cols, c1 + guard_width)

		guard_mask = np.zeros((self.rows, self.cols), dtype=bool)
		guard_mask[gr0:gr1, gc0:gc1] = True
		guard_mask[r0:r1, c0:c1] = False

		# Guard band: overdamped FHN params with zero diffusion.
		# Purpose: prevents wave activity from leaking into or out of the
		# storage region.  The guard cells sit at V_REST and cannot oscillate.
		self.a_map[guard_mask] = A_GUARD
		self.b_map[guard_mask] = B_GUARD
		self.eps_map[guard_mask] = EPS_GUARD
		self.d[:, guard_mask] = 0.0

		# Storage zone: bistable FHN params with zero diffusion.
		# A_STORE/B_STORE place the nullclines so that V_HIGH and V_LOW are
		# both stable fixed points (bistable regime).  Zero diffusion means
		# each cell holds its state independently — a voltage latch.
		self.a_map[r0:r1, c0:c1] = A_STORE
		self.b_map[r0:r1, c0:c1] = B_STORE
		self.eps_map[r0:r1, c0:c1] = EPS_STORE
		self.d[:, r0:r1, c0:c1] = 0.0

		# Init at LOW attractor
		self.v[r0:r1, c0:c1] = np.float32(V_LOW)
		self.w[r0:r1, c0:c1] = np.float32(W_LOW)
		# Guard cells at rest
		self.v[guard_mask] = np.float32(V_REST)
		self.w[guard_mask] = np.float32(W_REST)

		# Freeze storage + guard from sparse frontier
		self.storage_freeze_mask[gr0:gr1, gc0:gc1] = True

		self.param_zones[name] = {
			"r0": int(r0), "r1": int(r1), "c0": int(c0), "c1": int(c1),
			"zone_rows": int(zone_rows), "zone_cols": int(zone_cols),
			"guard_width": int(guard_width),
			"kind": "storage",
			"keys": [],
			"patch_rows": int(patch_rows) if patch_rows > 0 else int(zone_rows),
			"patch_cols": int(patch_cols) if patch_cols > 0 else int(zone_cols),
			"slots": {},
			"free_slots": [],
		}
		# Pre-compute slot grid
		zone = self.param_zones[name]
		pr, pc = zone["patch_rows"], zone["patch_cols"]
		slot_idx = 0
		sr = r0
		while sr + pr <= r1:
			sc = c0
			while sc + pc <= c1:
				zone["free_slots"].append(slot_idx)
				slot_idx += 1
				sc += pc
			sr += pr
		zone["total_slots"] = slot_idx
		return True

	def deallocate_storage_zone(self, name: str) -> bool:
		"""Deallocate a storage zone, restoring default parameters."""
		zone = self.param_zones.pop(name, None)
		if zone is None:
			return False
		r0, r1, c0, c1 = zone["r0"], zone["r1"], zone["c0"], zone["c1"]
		gw = zone.get("guard_width", 2)
		gr0, gr1 = max(0, r0 - gw), min(self.rows, r1 + gw)
		gc0, gc1 = max(0, c0 - gw), min(self.cols, c1 + gw)
		self.a_map[gr0:gr1, gc0:gc1] = A_FHN
		self.b_map[gr0:gr1, gc0:gc1] = B_FHN
		self.eps_map[gr0:gr1, gc0:gc1] = EPS
		self.d[:, gr0:gr1, gc0:gc1] = DIFFUSE
		self.v[gr0:gr1, gc0:gc1] = np.float32(V_REST)
		self.w[gr0:gr1, gc0:gc1] = np.float32(W_REST)
		self.storage_freeze_mask[gr0:gr1, gc0:gc1] = False
		# Clean up stored keys from store
		for k in zone.get("keys", []):
			if k in self.store.templates:
				del self.store.templates[k]
		return True

	def storage_zone_capacity(self, name: str) -> int:
		"""Return number of storage bits (cells) in a zone."""
		zone = self.param_zones.get(name)
		if zone is None or zone.get("kind") != "storage":
			return 0
		return zone["zone_rows"] * zone["zone_cols"]

	def _slot_bounds(self, zone: dict, slot_idx: int) -> Tuple[int, int, int, int]:
		"""Return (sr0, sr1, sc0, sc1) for a slot index within a zone."""
		r0, c0 = zone["r0"], zone["c0"]
		pr, pc = zone["patch_rows"], zone["patch_cols"]
		zone_cols_count = (zone["c1"] - c0) // pc
		slot_row = slot_idx // zone_cols_count
		slot_col = slot_idx % zone_cols_count
		sr0 = r0 + slot_row * pr
		sc0 = c0 + slot_col * pc
		return sr0, sr0 + pr, sc0, sc0 + pc

	# --- Grid I/O and content-addressable grep --------------------------------

	def save_to_grid(self, key: str, bits: np.ndarray, zone_name: str) -> bool:
		"""Write a bit pattern to the grid as bistable voltage latches.

		Each bit becomes one grid cell: 1 → V_HIGH (upper fixed point),
		0 → V_LOW (lower fixed point).  Recovery variable w is set to the
		matching attractor so the latch is immediately stable.

		If the zone has slots, allocates a slot for this key (or overwrites
		an existing slot for the same key).
		"""
		zone = self.param_zones.get(zone_name)
		if zone is None or zone.get("kind") != "storage":
			return False
		flat = np.asarray(bits, dtype=np.int32).ravel()
		pr, pc = zone["patch_rows"], zone["patch_cols"]
		slot_size = pr * pc

		# Find or allocate a slot
		slots = zone.get("slots", {})
		if key in slots:
			slot_idx = slots[key]
		elif zone.get("free_slots"):
			slot_idx = zone["free_slots"].pop(0)
			slots[key] = slot_idx
			zone["slots"] = slots
		else:
			# Single-slot legacy: use entire zone
			slot_idx = -1

		if slot_idx >= 0:
			sr0, sr1, sc0, sc1 = self._slot_bounds(zone, slot_idx)
		else:
			sr0, sr1 = zone["r0"], zone["r1"]
			sc0, sc1 = zone["c0"], zone["c1"]
		n_cells = (sr1 - sr0) * (sc1 - sc0)
		if flat.size < n_cells:
			flat = np.pad(flat, (0, n_cells - flat.size), constant_values=0)
		else:
			flat = flat[:n_cells]

		patch = flat.reshape(sr1 - sr0, sc1 - sc0)
		self.v[sr0:sr1, sc0:sc1] = np.where(patch, np.float32(V_HIGH), np.float32(V_LOW))
		self.w[sr0:sr1, sc0:sc1] = np.where(patch, np.float32(W_HIGH), np.float32(W_LOW))
		if key not in zone.get("keys", []):
			zone.setdefault("keys", []).append(key)
		return True

	def free_grid_slot(self, key: str, zone_name: str) -> bool:
		"""Free a slot for a key, returning it to the free list."""
		zone = self.param_zones.get(zone_name)
		if zone is None:
			return False
		slots = zone.get("slots", {})
		if key not in slots:
			return False
		slot_idx = slots.pop(key)
		sr0, sr1, sc0, sc1 = self._slot_bounds(zone, slot_idx)
		self.v[sr0:sr1, sc0:sc1] = np.float32(V_LOW)
		self.w[sr0:sr1, sc0:sc1] = np.float32(W_LOW)
		zone.get("free_slots", []).append(slot_idx)
		if key in zone.get("keys", []):
			zone["keys"].remove(key)
		return True

	def load_from_grid(self, key: str, zone_name: str) -> Optional[np.ndarray]:
		"""Read a stored bit pattern from the grid via threshold decode."""
		zone = self.param_zones.get(zone_name)
		if zone is None or zone.get("kind") != "storage":
			return None
		slots = zone.get("slots", {})
		if key in slots:
			sr0, sr1, sc0, sc1 = self._slot_bounds(zone, slots[key])
		else:
			sr0, sr1 = zone["r0"], zone["r1"]
			sc0, sc1 = zone["c0"], zone["c1"]
		v_stored = self.v[sr0:sr1, sc0:sc1].ravel()
		return (v_stored > np.float32(BISTABLE_THRESHOLD)).astype(np.int32)

	def grep_storage(self, needle_bits: np.ndarray, zone_names: Optional[List[str]] = None, max_results: int = 16, max_hd: int = -1) -> List[dict]:
		"""Content-addressable search via analytic flip-count probe across storage zones.
		
		If max_hd >= 0, only returns matches with hamming_distance <= max_hd.
		Use max_hd=0 for exact-match lookup.
		"""
		if zone_names is None:
			zone_names = [n for n, z in self.param_zones.items() if z.get("kind") == "storage"]
		results = []
		needle = np.asarray(needle_bits, dtype=np.int32).ravel()
		for zn in zone_names:
			zone = self.param_zones.get(zn)
			if zone is None or zone.get("kind") != "storage":
				continue
			slots = zone.get("slots", {})
			if slots:
				# Probe each slot independently
				pr, pc = zone["patch_rows"], zone["patch_cols"]
				slot_size = pr * pc
				if needle.size != slot_size:
					continue
				for key, slot_idx in slots.items():
					sr0, sr1, sc0, sc1 = self._slot_bounds(zone, slot_idx)
					stored_v = self.v[sr0:sr1, sc0:sc1].ravel()
					hd = self._probe_hamming(stored_v, needle)
					if max_hd >= 0 and hd > max_hd:
						continue
					results.append({"key": key, "zone": zn, "hamming_distance": hd})
			else:
				# Legacy single-pattern zone
				r0, r1, c0, c1 = zone["r0"], zone["r1"], zone["c0"], zone["c1"]
				n_cells = (r1 - r0) * (c1 - c0)
				if needle.size != n_cells:
					continue
				stored_v = self.v[r0:r1, c0:c1].ravel()
				hd = self._probe_hamming(stored_v, needle)
				if max_hd >= 0 and hd > max_hd:
					continue
				for k in zone.get("keys", []):
					results.append({"key": k, "zone": zn, "hamming_distance": hd})
		results.sort(key=lambda x: x["hamming_distance"])
		return results[:max_results]

	def grep_storage_exact(self, needle_bits: np.ndarray, zone_names: Optional[List[str]] = None) -> Optional[dict]:
		"""Exact pattern lookup (HD=0). Returns first match or None."""
		matches = self.grep_storage(needle_bits, zone_names=zone_names, max_results=1, max_hd=0)
		return matches[0] if matches else None

	def _probe_hamming(self, stored_v: np.ndarray, needle: np.ndarray) -> int:
		"""Analytic flip-count Hamming probe — no bit readback required.

		Injects a small perturbation (STORAGE_PROBE_AMP) aligned with the
		needle pattern.  For cells where stored and needle agree, the nudge
		pushes the voltage further into the same basin — no flip.  Where
		they disagree, the nudge may push the voltage across the bistable
		boundary — a flip.  The count of flips equals the Hamming distance.

		This runs at O(n) with a single vectorised pass, giving flat-cost
		content-addressable search across all stored patterns.
		"""
		probe_delta = np.float32(STORAGE_PROBE_AMP) * (np.float32(2.0) * needle.astype(np.float32) - np.float32(1.0))
		v_probed = stored_v + probe_delta
		stored_high = stored_v > np.float32(BISTABLE_THRESHOLD)
		probed_high = np.where(stored_high,
			v_probed >= np.float32(V_BOUNDARY_HIGH),
			v_probed > np.float32(V_BOUNDARY_LOW))
		return int(np.sum(stored_high != probed_high))

	# --- Sprite blit (display compositing primitives) -------------------------

	def _load_sprite(self, zone_name: str, key: str) -> Optional[Tuple[np.ndarray, int, int]]:
		"""Load a stored bit-pattern as a 2D sprite array.

		Returns (sprite_2d, patch_rows, patch_cols) or None if the key
		doesn't exist or the zone geometry is invalid.
		"""
		zone = self.param_zones.get(zone_name)
		if zone is None or zone.get("kind") != "storage":
			return None
		bits = self.load_from_grid(key, zone_name)
		if bits is None:
			return None
		pr, pc = zone.get("patch_rows", 0), zone.get("patch_cols", 0)
		if pr <= 0 or pc <= 0 or bits.size != pr * pc:
			return None
		return bits.reshape(pr, pc), pr, pc

	def _hw_sprite_blit(
		self,
		zone_name: str,
		key: str,
		dst_r0: int,
		dst_c0: int,
		freeze: bool = True,
		on_value: float = float(V_HIGH),
		off_value: float = float(V_LOW),
	) -> bool:
		"""Blit a stored sprite onto the voltage grid."""
		self._bump_syscall("sprite_blit")
		result = self._load_sprite(zone_name, key)
		if result is None:
			return False
		sprite, pr, pc = result
		dr0, dc0 = int(dst_r0), int(dst_c0)
		dr1, dc1 = dr0 + pr, dc0 + pc
		if dr0 < 0 or dc0 < 0 or dr1 > self.rows or dc1 > self.cols:
			return False
		sprite_bool = sprite.astype(bool)
		self.v[dr0:dr1, dc0:dc1] = np.where(sprite_bool, np.float32(on_value), np.float32(off_value))
		if freeze:
			self.storage_freeze_mask[dr0:dr1, dc0:dc1] = True
		return True

	def _sprite_stamp_to_buffer(
		self,
		target: "np.ndarray",
		zone_name: str,
		key: str,
		dst_r0: int,
		dst_c0: int,
		on_value: float = float(V_HIGH),
		off_value: float = float(V_LOW),
		transparent_off: bool = True,
	) -> bool:
		"""Paint a stored sprite onto an arbitrary 2D float buffer."""
		self._bump_syscall("sprite_stamp_to_buffer")
		result = self._load_sprite(zone_name, key)
		if result is None:
			return False
		sprite, pr, pc = result
		if not isinstance(target, np.ndarray) or target.ndim < 2:
			return False
		th, tw = target.shape[0], target.shape[1]
		dr0, dc0 = int(dst_r0), int(dst_c0)
		dr1, dc1 = dr0 + pr, dc0 + pc
		if dr0 < 0 or dc0 < 0 or dr1 > th or dc1 > tw:
			return False
		sprite_bool = sprite.astype(bool)
		if transparent_off:
			target[dr0:dr1, dc0:dc1][sprite_bool] = np.float32(on_value)
		else:
			target[dr0:dr1, dc0:dc1] = np.where(sprite_bool, np.float32(on_value), np.float32(off_value))
		return True

	# --- Zone introspection ----------------------------------------------------

	def list_storage_zones(self) -> List[dict]:
		"""List all allocated storage zones and their metadata."""
		out = []
		for name, zone in self.param_zones.items():
			if zone.get("kind") != "storage":
				continue
			out.append({
				"name": name,
				"r0": zone["r0"], "r1": zone["r1"],
				"c0": zone["c0"], "c1": zone["c1"],
				"capacity_bits": int(zone["zone_rows"]) * int(zone["zone_cols"]),
				"keys": list(zone.get("keys", [])),
			})
		return out