#!/usr/bin/env python3
"""
Hyperdimensional Vocab Memory — oscillator associative memory for a language model.

  - Encode vocab as 256 complex oscillator hypervectors (random projection).
  - Burn in token co-occurrence from text corpus (Hebbian relationship modulation).
  - Query: context string → recall associated tokens by phase-coherent pattern
    completion.
  - Fuse: memory scores boost model logits (3B params for reasoning, memory for
    knowledge → functions like a larger model).

Real Llama 3B tokenizer + bf16 embeddings, text corpus from the local disk.

    python3 hyper_vocab_memory.py
"""
import glob, json, math, os, struct, sys, time
import numpy as np

# ── tokenizer + embeddings (llama 3B, mmap'd bf16) ──────────────────────
BASE = glob.glob("/home/compunerd/.cache/huggingface/hub/"
                 "models--huihui-ai--Hermes-3-Llama-3.2-3B-abliterated/"
                 "snapshots/*")[0]
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                "bqsm_assist"))
# use the Safetensors helper inline rather than import (avoids directory issues)
C = json.load(open(os.path.join(BASE, "config.json")))
D = C["hidden_size"]
TK = json.load(open(os.path.join(BASE, "tokenizer.json")))
VOCAB = TK["model"]["vocab"]    # token -> id
INV = {v: k for k, v in VOCAB.items()}


class ST:
    """Minimal safetensors reader — mmap, no copy on parse."""
    def __init__(self):
        self.shards, self.idx = [], {}
        for p in sorted(glob.glob(os.path.join(BASE, "*.safetensors"))):
            mm = np.memmap(p, dtype=np.uint8, mode="r")
            n = struct.unpack("<Q", bytes(mm[:8]))[0]
            hdr = json.loads(bytes(mm[8:8 + n]).decode())
            st = 8 + n
            si = len(self.shards)
            self.shards.append((mm, st))
            for k, v in hdr.items():
                if k != "__metadata__":
                    self.idx[k] = (si, v)

    def get(self, name):
        si, v = self.idx[name]
        mm, st = self.shards[si]
        a, b = v["data_offsets"]
        raw = np.asarray(mm[st + a: st + b])
        dt = v["dtype"]
        if dt == "BF16":
            return ((raw.view(np.uint16).astype(np.uint32) << 16)
                    ).view(np.float32).reshape(v["shape"])
        elif dt == "F16":
            return raw.view(np.float16).astype(np.float32).reshape(v["shape"])
        return raw.view(np.float32).reshape(v["shape"])


st = ST()


# ── embed-row helper (bf16 zero-copy per row) ───────────────────────────
def embed_row(token_id):
    si, v = st.idx["model.embed_tokens.weight"]
    mm, st_off = st.shards[si]
    D_emb = v["shape"][1]
    a, b = v["data_offsets"]
    off = st_off + a + token_id * D_emb * 2  # bf16 = 2 bytes
    raw = np.asarray(mm[off: off + D_emb * 2])
    return ((raw.view(np.uint16).astype(np.uint32) << 16)
            ).view(np.float32)  # → f32, D-dim


# ── tokenizer ────────────────────────────────────────────────────────────


def encode(text):
    ids, words = [128000], text.split()
    for i, w in enumerate(words):
        w = w.strip(".,;:!?'\"()[]{}<>/\\|+-=*&^%$#@~`").strip()
        if not w:
            continue
        key = ("Ġ" + w) if i else w
        if key in VOCAB:
            ids.append(VOCAB[key])
        elif w in VOCAB:
            ids.append(VOCAB[w])
        else:
            for ch in key:
                if ch in VOCAB:
                    ids.append(VOCAB[ch])
    return ids


def dec(i):
    return INV.get(i, f"[{i}]").replace("Ġ", " ").replace("Ċ", "\n")


# ── Hyperdimensional encoding: token embedding → oscillator state ───────
N_OSC = 256
rng = np.random.default_rng(42)
# random projection matrix [2*N_OSC, D]
PROJ = rng.standard_normal((N_OSC * 2, D), dtype=np.float32) / np.sqrt(D)


def osc_vector(tok, cache):
    """Complex oscillator state for a token (from cache), or zeros if unseen."""
    v = cache.get(tok)
    return v if v is not None else np.zeros(N_OSC, dtype=np.complex64)


def build_cache(token_ids, chunk=4096, label=""):
    """Encode token IDs into normalized complex oscillator states (batched).

    Reads the bf16 embedding matrix via mmap in chunks, projects each chunk,
    and stores a compact complex64 vector per token.  Full vocab (~128K) is
    ~262 MB of cache — built chunk-wise so peak memory stays ~100 MB."""
    si, v = st.idx["model.embed_tokens.weight"]
    mm, st_off = st.shards[si]
    D_emb = v["shape"][1]
    a, b = v["data_offsets"]
    base = st_off + a
    ids = list(token_ids)
    cache = {}
    for i in range(0, len(ids), chunk):
        batch = ids[i:i + chunk]
        lo, hi = batch[0], batch[-1]
        off = base + lo * D_emb * 2
        raw = np.asarray(mm[off: off + (hi - lo + 1) * D_emb * 2])
        E = ((raw.view(np.uint16).astype(np.uint32) << 16)
             ).view(np.float32).reshape(hi - lo + 1, D_emb)
        P = E[[t - lo for t in batch]] @ PROJ.T   # [batch, 2*N_OSC]
        P /= (np.linalg.norm(P, axis=1, keepdims=True) + 1e-8)
        C = (P[:, :N_OSC] + 1j * P[:, N_OSC:]).astype(np.complex64)
        for j, t in enumerate(batch):
            cache[t] = C[j]
        if label and (i == 0 or (i + chunk) >= len(ids)):
            print(f"    {label} {min(i + chunk, len(ids))}/{len(ids)}", flush=True)
    return cache


# ── Gather the tokens that actually matter (corpus + queries) ───────────
corpus_files = [
    "/home/compunerd/agent_framework/README.md",
    "/home/compunerd/agent_framework/bqsm_assist/WAVE_RIDER_BREAKTHROUGH.md",
    "/home/compunerd/Desktop/bqsm/basin-quotient-machine/LENS_CONTROL_METHODS.md",
    "/home/compunerd/Desktop/bqsm/basin-quotient-machine/README.md",
]
test_queries = [
    "The capital of France is",
    "BQSM uses coupled",
    "The ring computes through mode",
    "lens site 0 enhances the",
    "Phase 0 Gate",
    "a transformer forward pass as",
    "the model with real bf16",
    "attention becomes geometric",
]

used_ids = set()
corpus_texts = []
for fp in corpus_files:
    if os.path.exists(fp):
        text = open(fp).read()[:50000]
        corpus_texts.append(text)
        used_ids.update(encode(text))
for q in test_queries:
    used_ids.update(encode(q))

# Full-dictionary coverage: encode the entire vocab so ANY token the model
# emits can be scored/recalled.  Corpus-derived associations still come from
# `used_ids`, but every token now has a hypervector.  BQSM_FULL_VOCAB=off
# reverts to the corpus-subset cache for fast smoke tests.
FULL_VOCAB = os.environ.get("BQSM_FULL_VOCAB", "on") != "off"
if FULL_VOCAB:
    V = st.idx["model.embed_tokens.weight"][1]["shape"][0]
    print(f"encoding full vocab ({V} tokens) — one-time, ~30s...")
    t0 = time.time()
    cache = build_cache(range(V), label="vocab")
    print(f"  {time.time()-t0:.1f}s ({len(cache)} tokens cached)")
else:
    print(f"encoding {len(used_ids)} distinct tokens (corpus + queries)...")
    t0 = time.time()
    cache = build_cache(sorted(used_ids))
    print(f"  {time.time()-t0:.1f}s")

# Burn-in + build sparse "following" index (skip-gram, distance-decayed)
from collections import Counter, defaultdict

print("\nBurn-in corpus (skip-gram PMI, window=3)...")
MAX_DIST = 3
unigram = Counter()
skipgram = {d: Counter() for d in range(1, MAX_DIST + 1)}
total_tokens = 0
for text in corpus_texts:
    ids = encode(text)
    total_tokens += len(ids)
    unigram.update(ids)
    for d in range(1, MAX_DIST + 1):
        skipgram[d].update(zip(ids[:-d], ids[d:]))

# Distance-decayed PMI: tokens d apart get weight 1/d.  This captures
# "France -> is -> Paris" as "France -> Paris" (d=2, weight 0.5), which is
# what a pure bigram memory misses.
W = np.zeros((N_OSC, N_OSC), dtype=np.complex64)
following = defaultdict(list)       # token_id -> [(target_id, weight), ...]
n_pairs = 0
for d in range(1, MAX_DIST + 1):
    decay = 1.0 / d
    for (a, b), cnt in skipgram[d].items():
        za = cache.get(a); zb = cache.get(b)
        if za is None or zb is None:
            continue
        pmi = math.log((cnt * total_tokens) / (unigram[a] * unigram[b]) + 1e-12)
        if pmi <= 0:
            continue
        w = decay * pmi
        W += w * np.outer(za, np.conj(zb))
        following[a].append((b, w))
        n_pairs += 1
print(f"  {n_pairs} associations burned in (PMI>0), from {total_tokens} tokens")
norm = np.linalg.norm(W)
if norm > 0:
    W /= norm

# ── Sparse recall: per context token, aggregate its strongest followers ───
def query_sparse(context_str, top_k=20):
    ids = encode(context_str)
    scores = defaultdict(float)
    for i, cid in enumerate(ids):
        if cid not in following:
            continue
        # last token gets 2× weight for next-token prediction
        w = 2.0 if i == len(ids) - 1 else 1.0
        for tid, pmi in following[cid]:
            scores[tid] += w * pmi
    ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    return [r for r in ranked if r[0] in cache][:top_k]

# ── Query: context → associative recall ──────────────────────────────────
def query(context_str, top_k=20):
    """Encode context as oscillator state, recall associated tokens via W."""
    ids = encode(context_str)
    ctx = np.zeros(N_OSC, dtype=np.complex64)
    n_valid = 0
    for t in ids:
        z = cache.get(t)
        if z is not None:
            ctx += z
            n_valid += 1
    if n_valid == 0:
        return []
    ctx /= n_valid + 1e-8
    recalled = W @ ctx
    recalled = recalled / (np.linalg.norm(recalled) + 1e-8)
    scores = []
    for t, zt in cache.items():
        score = float(abs(np.dot(np.conj(zt), recalled)))
        scores.append((t, score))
    scores.sort(key=lambda x: x[1], reverse=True)
    return scores[:top_k]


# ── Demo (only when run directly) ─────────────────────────────────────────
if __name__ == "__main__":
    print("\n" + "=" * 66)
    print("HYPERDIMENSIONAL VOCAB MEMORY — recall demo")
    print("=" * 66)

    tests = [
        ("The capital of France is", "Paris"),
        ("BQSM uses coupled", "oscillator"),
        ("The ring computes through mode", "coupling"),
        ("lens site 0 enhances the", "channel"),
        ("Phase 0 Gate", "FAILURE"),
        ("a transformer forward pass as", "coupled"),
        ("the model with real bf16", "weights"),
        ("attention becomes geometric", "adjacency"),
    ]

    def find_token(text):
        for t in cache:
            if dec(t).strip() == text:
                return t
        return None

    for context, expected in tests:
        results = query_sparse(context)
        expected_id = find_token(expected)
        rank = None
        for i, (t, s) in enumerate(results):
            if t == expected_id:
                rank = i + 1
                break
        print(f"\n  \"{context}\"")
        print(f"  expect: \"{expected}\"  rank: "
              f"{rank if rank else '-- (not in top %d)' % len(results)}")
        print(f"  top 5: ", end="")
        for t, s in results[:5]:
            print(f"{dec(t)!r}({s:.4f})", end="  ")
        print()

    print("\n" + "=" * 66)
    print("FUSION — memory boosts model logits (simulated)")
    print("=" * 66)
    context = "The capital of France is"
    model_logits = {t: float(rng.standard_normal()) * 0.5 for t in cache}
    results = query_sparse(context)
    for t, mem_score in results:
        model_logits[t] = model_logits.get(t, 0.0) + 2.0 * mem_score
    top_after = sorted(model_logits, key=lambda t: model_logits[t],
                       reverse=True)[:10]
    print(f"  context: {context!r}")
    print(f"  top-10 after fusion: {[dec(t) for t in top_after]}")
    paris_id = find_token("Paris")
    if paris_id:
        rank = top_after.index(paris_id) + 1 if paris_id in top_after else None
        print(f"  'Paris' rank after fusion: "
              f"{'#' + str(rank) if rank else '-- (out of top 10)'}")

    print("\n" + "=" * 66)
    print("HOW IT SCALES TO 30B-CLASS:")
    print("  - 3B model: grammar, reasoning, common patterns (its parameters)")
    print("  - Oscillator memory: facts, entity links, co-occurrence (burn-in)")
    print("  - The memory costs N² oscillators (~256² = 65K couplings), not GBs")
    print("  - Continually learns: new facts burn in without retraining the model")
    print("  - Hyperdimensional encoding: near-orthogonal random projections")
    print("    = associative memory for 128K vocab in ~65K complex couplings")
    print("=" * 66)