code/encompassment_engine_v2.py

"""
Encompassment Engine V2 — K6 as CEO (not as counter).

Fixes V1's fatal flaw: V1's centrality was 80% K1+K3+K5 duplication.
V2 uses ONLY new information:
  - spread: how many cells share this pattern
  - diversity: how many nearby cells (from plexus)
  - uniqueness: how isolated this pattern is

V2 DECIDES (not measures):
  - decide_split: K6 says yes/no (with S as secondary)
  - decide_birth: K6 says yes/no
  - guide_retrieval: K6 ranks cells for attention
"""
from collections import defaultdict
from typing import Dict, List, Optional
import math


class EncompassmentEngineV2:
    def __init__(self, max_index_size=10000):
        self.pattern_index: Dict[str, List[str]] = defaultdict(list)
        self.centrality: Dict[str, float] = {}
        self.cell_metadata: Dict[str, dict] = {}
        self.birth_history: List[dict] = []
        self.split_history: List[dict] = []
        self.max_index_size = max_index_size
        self._total_registrations = 0

    # ------------------------------------------------------------------
    # Registration
    # ------------------------------------------------------------------
    def register(self, cell_id: str, fp: str, layer_idx: int = -1):
        """Register a cell's pattern fingerprint. K1-K5 NOT used."""
        if not fp:
            return
        if cell_id in self.cell_metadata:
            old_fp = self.cell_metadata[cell_id].get("fp")
            if old_fp == fp:
                self._recompute_centrality(cell_id)
                return
            if old_fp and cell_id in self.pattern_index.get(old_fp, []):
                self.pattern_index[old_fp].remove(cell_id)

        self.pattern_index[fp].append(cell_id)
        self.cell_metadata[cell_id] = {"fp": fp, "layer_idx": layer_idx}
        self._recompute_centrality(cell_id)
        self._total_registrations += 1

        if len(self.pattern_index) > self.max_index_size:
            self._prune()

    def _count_nearby_cells(self, cell_id: str, hamming_fn, threshold: float = 0.6) -> int:
        """Count cells whose pattern is *close but not identical* to this one."""
        meta = self.cell_metadata.get(cell_id)
        if not meta:
            return 0
        my_fp = meta["fp"]
        count = 0
        for fp, members in self.pattern_index.items():
            if fp == my_fp:
                continue
            try:
                sim = hamming_fn(my_fp, fp)
            except Exception:
                continue
            if sim > threshold:
                count += len(members)
        return count

    def _recompute_centrality(self, cell_id: str,
                               hamming_fn=None, threshold: float = 0.6):
        """
        K6 = new information ONLY:
          - spread: how many cells share this exact pattern
          - diversity: how many nearby (similar but not identical) cells
        NO K1, K3, K5.
        """
        meta = self.cell_metadata.get(cell_id)
        if not meta:
            self.centrality[cell_id] = 0.0
            return

        fp = meta["fp"]
        n_sharing = len(self.pattern_index.get(fp, []))

        # spread: [0, 1]
        spread = min(1.0, math.log1p(n_sharing) / math.log1p(20))

        # diversity: size of overall pattern space (how many distinct patterns)
        n_patterns = len(self.pattern_index)
        diversity = min(1.0, math.log1p(n_patterns) / math.log1p(50))

        # centrality = CEO score (pure network info)
        c = 0.5 * spread + 0.5 * diversity
        self.centrality[cell_id] = max(0.0, min(1.0, c))

    def refresh_all(self, hamming_fn=None):
        """Recompute all centralities (call after several registrations)."""
        for cid in list(self.cell_metadata.keys()):
            self._recompute_centrality(cid, hamming_fn=hamming_fn)

    def _prune(self):
        keep = {}
        for fp, members in self.pattern_index.items():
            if len(members) >= 2:
                keep[fp] = members
            else:
                for m in members:
                    if self.centrality.get(m, 0.0) > 0.5:
                        keep[fp] = members
                        break
        self.pattern_index = defaultdict(list, keep)

    # ------------------------------------------------------------------
    # K6 read
    # ------------------------------------------------------------------
    def K6(self, cell_id: str) -> float:
        return float(self.centrality.get(cell_id, 0.0))

    def K6_array(self, cell_ids: List[str]) -> List[float]:
        return [self.K6(c) for c in cell_ids]

    # ------------------------------------------------------------------
    # CEO decisions (not measurements)
    # ------------------------------------------------------------------
    def decide_split(self, cell_id: str, S_value: float, S_star: float) -> bool:
        """
        Split decision:
          - S >= S* (saturation signal)
          - AND centrality >= 0.5 (pattern is in active region of graph)
        Centrality prevents wasting splits on isolated noise patterns.
        """
        c = self.centrality.get(cell_id, 0.0)
        return (S_value >= S_star) and (c >= 0.15)   # threshold for R7 scale

    def decide_birth(self, fp: str, hamming_fn, threshold: float = 0.55) -> bool:
        """
        Birth = pattern that is truly new.
        - Not already in index
        - Not similar to any existing pattern
        """
        if not fp:
            return False
        if fp in self.pattern_index and len(self.pattern_index[fp]) > 0:
            return False
        for existing_fp in self.pattern_index:
            try:
                sim = hamming_fn(fp, existing_fp)
            except Exception:
                continue
            if sim > threshold:
                return False
        return True

    def record_birth(self, cell_id: str, parent_id: str, fp: str, step: int):
        self.birth_history.append({
            "cell_id": cell_id, "parent_id": parent_id,
            "fp": fp, "step": step,
        })

    def record_split(self, parent_id: str, children: List[str],
                     S: float, K6: float, step: int):
        self.split_history.append({
            "parent_id": parent_id, "children": children,
            "S": round(float(S), 4), "K6": round(float(K6), 4),
            "step": step,
        })

    # ------------------------------------------------------------------
    # Retrieval guidance
    # ------------------------------------------------------------------
    def guide_retrieval(self, query_fp: str, hamming_fn,
                         top_k: int = 5) -> List[str]:
        if not query_fp:
            return []
        scores = []
        for fp, members in self.pattern_index.items():
            try:
                sim = hamming_fn(query_fp, fp)
            except Exception:
                sim = 0.0
            if sim < 0.1:
                continue
            for cid in members:
                c = self.centrality.get(cid, 0.0)
                scores.append((sim * (0.5 + 0.5 * c), cid))
        scores.sort(reverse=True)
        return [cid for _, cid in scores[:top_k]]

    # ------------------------------------------------------------------
    # Diagnostics
    # ------------------------------------------------------------------
    def summary(self) -> dict:
        n_patterns = len(self.pattern_index)
        n_cells = len(self.cell_metadata)
        if self.centrality:
            c_vals = list(self.centrality.values())
            c_min, c_max = min(c_vals), max(c_vals)
            c_mean = sum(c_vals) / len(c_vals)
        else:
            c_min = c_max = c_mean = 0.0
        return {
            "n_patterns": n_patterns,
            "n_cells_indexed": n_cells,
            "centrality_min": round(c_min, 4),
            "centrality_max": round(c_max, 4),
            "centrality_mean": round(c_mean, 4),
            "births_total": len(self.birth_history),
            "splits_total": len(self.split_history),
            "total_registrations": self._total_registrations,
        }


if __name__ == "__main__":
    def fake_hamming(a, b):
        if not a or not b:
            return 0.0
        n = min(len(a), len(b))
        return sum(1 for i in range(n) if a[i] == b[i]) / n

    e = EncompassmentEngineV2()
    # Unique pattern
    e.register("cell-1", "aaaaaaaaaa", layer_idx=0)
    # Same pattern (spread)
    e.register("cell-2", "aaaaaaaaaa", layer_idx=0)
    e.register("cell-3", "aaaaaaaaaa", layer_idx=0)
    # Similar but not identical (diversity)
    e.register("cell-4", "aaaaaaaabb", layer_idx=0)
    # Far away
    e.register("cell-5", "ffffffffff", layer_idx=0)

    e.refresh_all(hamming_fn=fake_hamming)

    print("=== K6 V2 (CEO) ===")
    print(f"cell-1 (with 2 sharing + nearby): K6 = {e.K6('cell-1'):.4f}")
    print(f"cell-5 (isolated): K6 = {e.K6('cell-5'):.4f}")
    print(f"summary: {e.summary()}")
    print()
    print("=== Decisions ===")
    print(f"decide_birth(new pattern): {e.decide_birth('zzzzzzzzzz', fake_hamming)}")
    print(f"decide_birth(existing): {e.decide_birth('aaaaaaaaaa', fake_hamming)}")
    print(f"decide_split(cell-1, S=0.9, S*=0.85): {e.decide_split('cell-1', 0.9, 0.85)}")
    print(f"decide_split(cell-5, S=0.9, S*=0.85): {e.decide_split('cell-5', 0.9, 0.85)}")