#!/usr/bin/env python3
"""
Vector Engine — Core 24-bit Substrate Operations
=================================================
Maps words, tokens, and patches to 24-bit Golay codewords.
Provides NRCI scoring, Hamming distance, and lattice classification.
"""
import sys, os, hashlib, re, json
from typing import List, Dict, Tuple, Optional, Set
from fractions import Fraction
from collections import Counter

# Add parent dir for UBP engine
_this_dir = os.path.dirname(os.path.abspath(__file__))
_parent = os.path.dirname(_this_dir)
if _parent not in sys.path:
    sys.path.insert(0, _parent)

from ubp_unified_v5 import (
    GolayCodeEngine, LeechLatticeEngine, UBPSourceCodeParticlePhysics,
)

# ── Singletons (expensive to create, reuse forever) ──
_GOLAY: Optional[GolayCodeEngine] = None
_LEECH: Optional[LeechLatticeEngine] = None
_PP: Optional[UBPSourceCodeParticlePhysics] = None

def get_golay() -> GolayCodeEngine:
    global _GOLAY
    if _GOLAY is None:
        _GOLAY = GolayCodeEngine()
    return _GOLAY

def get_leech() -> LeechLatticeEngine:
    global _LEECH
    if _LEECH is None:
        g = get_golay()
        _LEECH = LeechLatticeEngine(g)
    return _LEECH

def get_pp() -> UBPSourceCodeParticlePhysics:
    global _PP
    if _PP is None:
        _PP = UBPSourceCodeParticlePhysics()
    return _PP


# ══════════════════════════════════════════════════════════════════════════════
# VECTOR MAPPING
# ══════════════════════════════════════════════════════════════════════════════

def word_to_hash24(word: str) -> List[int]:
    """Hash a word to 24 bits via SHA-256 prefix. Deterministic."""
    h = hashlib.sha256(word.lower().strip().encode()).digest()
    return [(h[i // 8] >> (7 - i % 8)) & 1 for i in range(24)]

def snap_to_codeword(vec: List[int]) -> Tuple[List[int], dict]:
    """Snap any 24-bit vector to nearest Golay codeword."""
    g = get_golay()
    return g.snap_to_codeword(list(vec))

def hamming(a: List[int], b: List[int]) -> int:
    """Hamming distance between two 24-bit vectors."""
    return sum(x ^ y for x, y in zip(a, b))

def vector_to_hex(vec: List[int]) -> int:
    """Convert 24-bit vector to integer."""
    return sum((1 << (23 - i)) for i in range(24) if vec[i])

def hex_to_vector(h: int) -> List[int]:
    """Convert integer to 24-bit vector."""
    return [(h >> (23 - i)) & 1 for i in range(24)]


# ══════════════════════════════════════════════════════════════════════════════
# LATTICE CLASSIFICATION
# ══════════════════════════════════════════════════════════════════════════════

LATTICE_CLASSES = {
    0: "Identity", 8: "Octad", 10: "Decad", 12: "Dodecad",
    14: "Tetradecad", 16: "Hexadecad", 20: "Icosad", 24: "Edge"
}

def lattice_class(hw: int) -> str:
    return LATTICE_CLASSES.get(hw, f"HW={hw}")

def classify(vec: List[int]) -> dict:
    """Full classification of a 24-bit vector."""
    l = get_leech()
    hw = sum(vec)
    nrci = float(l.calculate_nrci(vec))
    tax = float(l.calculate_symmetry_tax(vec))
    return {
        "hw": hw,
        "nrci": nrci,
        "tax": tax,
        "lattice": lattice_class(hw),
        "hex": f"0x{vector_to_hex(vec):06X}",
        "in_band": nrci >= 0.70,
        "on_octad": hw == 8,
    }


# ══════════════════════════════════════════════════════════════════════════════
# SVD VOCABULARY BUILDER
# ══════════════════════════════════════════════════════════════════════════════

class SVDVocabulary:
    """Builds distributional 24-bit vectors from a corpus using PPMI + SVD.
    Closely follows GLM20_svd_vocab.py methodology."""
    
    def __init__(self):
        self.word_vectors: Dict[str, List[int]] = {}
        self.word_snapped: Dict[str, List[int]] = {}
        self.word_meta: Dict[str, dict] = {}
        self.context_words: List[str] = []
        self._built = False
    
    def build_from_definitions(self, definitions: Dict[str, str],
                                context_size: int = 100,
                                window: int = 8,
                                n_dims: int = 24) -> int:
        """Build SVD vocabulary from word→definition mappings.
        Returns number of words mapped."""
        try:
            import numpy as np
        except ImportError:
            return 0
        
        # Tokenize corpus
        tokens = []
        for defn in definitions.values():
            tokens.extend(re.findall(r"[a-z]+", defn.lower()))
        tokens = [t for t in tokens if len(t) >= 3]
        
        target_words = sorted(definitions.keys())
        vocab_idx = {w: i for i, w in enumerate(target_words)}
        
        # Context words (most frequent non-target)
        freq = Counter(tokens)
        self.context_words = [w for w, _ in freq.most_common(context_size + len(target_words))
                              if w not in vocab_idx][:context_size]
        ctx_idx = {w: i for i, w in enumerate(self.context_words)}
        
        # PPMI co-occurrence
        cooc = np.zeros((len(target_words), len(self.context_words)))
        all_tokens = re.findall(r"[a-z]+", " ".join(definitions.values()).lower())
        for i, tok in enumerate(all_tokens):
            if tok not in vocab_idx:
                continue
            wi = vocab_idx[tok]
            for j in range(max(0, i - window), min(len(all_tokens), i + window + 1)):
                if j == i:
                    continue
                ctx = all_tokens[j]
                if ctx in ctx_idx:
                    cooc[wi, ctx_idx[ctx]] += 1
        
        # PPMI transform
        total = cooc.sum()
        if total == 0:
            return 0
        row_sums = cooc.sum(axis=1, keepdims=True)
        col_sums = cooc.sum(axis=0, keepdims=True)
        row_sums[row_sums == 0] = 1
        col_sums[col_sums == 0] = 1
        ppmi = np.log2((cooc * total) / (row_sums * col_sums) + 1e-10)
        ppmi[ppmi < 0] = 0
        
        # SVD → n_dims
        U, S, Vt = np.linalg.svd(ppmi, full_matrices=False)
        svd_vecs = U[:, :n_dims] * S[:n_dims]
        
        # Median-quantize to bits
        medians = np.median(svd_vecs, axis=0)
        bit_vecs = (svd_vecs > medians).astype(int)
        
        # Snap to Golay
        g = get_golay()
        l = get_leech()
        for i, word in enumerate(target_words):
            raw = [int(b) for b in bit_vecs[i]]
            self.word_vectors[word] = raw
            snapped, meta = g.snap_to_codeword(raw)
            self.word_snapped[word] = snapped
            self.word_meta[word] = {
                "raw_hw": sum(raw),
                "snapped_hw": sum(snapped),
                "nrci": float(l.calculate_nrci(snapped)),
                "lattice": lattice_class(sum(snapped)),
                "hex": f"0x{vector_to_hex(snapped):06X}",
                "method": "svd",
            }
        
        self._built = True
        return len(self.word_snapped)
    
    def get_vector(self, word: str) -> List[int]:
        """Get 24-bit vector for a word. Falls back to hash if not in SVD vocab."""
        w = word.lower().strip()
        if w in self.word_snapped:
            return self.word_snapped[w]
        # Fallback: hash + snap
        raw = word_to_hash24(w)
        snapped, _ = snap_to_codeword(raw)
        return snapped
    
    def get_meta(self, word: str) -> dict:
        """Get metadata for a word."""
        w = word.lower().strip()
        if w in self.word_meta:
            return self.word_meta[w]
        raw = word_to_hash24(w)
        snapped, _ = snap_to_codeword(raw)
        l = get_leech()
        return {
            "raw_hw": sum(raw),
            "snapped_hw": sum(snapped),
            "nrci": float(l.calculate_nrci(snapped)),
            "lattice": lattice_class(sum(snapped)),
            "hex": f"0x{vector_to_hex(snapped):06X}",
            "method": "hash_fallback",
        }
    
    def hamming_between(self, word_a: str, word_b: str) -> int:
        return hamming(self.get_vector(word_a), self.get_vector(word_b))
    
    def save(self, path: str):
        """Save vocabulary to JSON."""
        data = {
            "word_snapped": self.word_snapped,
            "word_meta": self.word_meta,
            "context_words": self.context_words,
        }
        with open(path, "w") as f:
            json.dump(data, f)
    
    def load(self, path: str) -> bool:
        """Load vocabulary from JSON."""
        if not os.path.exists(path):
            return False
        with open(path) as f:
            data = json.load(f)
        self.word_snapped = {k: list(v) for k, v in data["word_snapped"].items()}
        self.word_meta = data.get("word_meta", {})
        self.context_words = data.get("context_words", [])
        self._built = True
        return True


# ══════════════════════════════════════════════════════════════════════════════
# ZONE CENTROID (Idea Zone from GLM07)
# ══════════════════════════════════════════════════════════════════════════════

class IdeaZone:
    """Maintains a running EMA centroid for the current conversation topic."""
    
    def __init__(self, alpha: float = 0.3):
        self.alpha = alpha
        self.centroid: List[float] = [0.0] * 24
        self.words: List[str] = []
        self._snapped: Optional[List[int]] = None
    
    def update(self, word: str, vocab: SVDVocabulary):
        """Add a word to the zone, updating the EMA centroid."""
        vec = vocab.get_vector(word)
        self.centroid = [self.alpha * v + (1 - self.alpha) * c
                         for v, c in zip(vec, self.centroid)]
        self.words.append(word.lower())
        self._snapped = None  # invalidate cache
    
    def get_centroid(self) -> List[int]:
        """Get the snapped centroid vector."""
        if self._snapped is None:
            bits = [1 if c > 0.5 else 0 for c in self.centroid]
            self._snapped, _ = snap_to_codeword(bits)
        return self._snapped
    
    def get_centroid_nrci(self) -> float:
        l = get_leech()
        return float(l.calculate_nrci(self.get_centroid()))
    
    def distance_to(self, vec: List[int]) -> int:
        return hamming(self.get_centroid(), vec)
    
    def reset(self):
        self.centroid = [0.0] * 24
        self.words = []
        self._snapped = None


# ══════════════════════════════════════════════════════════════════════════════
# MOG (Miracle Octad Generator) Layout
# ══════════════════════════════════════════════════════════════════════════════

MOG_POSITIONS = [(r, c) for r in range(4) for c in range(6)]  # 24 positions

def mog_quadrants(vec: List[int]) -> List[int]:
    """Split 24-bit vector into 4 sextets (Reality/Information/Activation/Potential)."""
    return [sum(vec[i:i+6]) for i in range(0, 24, 6)]

def mog_dominant_layer(vec: List[int]) -> Tuple[str, int]:
    """Return the dominant ontological layer and its bit count."""
    layers = ["Reality", "Information", "Activation", "Potential"]
    q = mog_quadrants(vec)
    idx = q.index(max(q))
    return layers[idx], q[idx]