"""
revision_memory.py — belief revision with derivation-tracked cascade.

Fixes dream-cascade loose ends:
  - Derived edges have weight = min over all derivation paths
  - Removing a source triggers full re-derivation
  - Every derived edge knows its parents (explainability)
  - Chain quality = min substrate sim along the walk

Two views of the same memory:
  - The substrate trace (what the phasor memory can answer on its own)
  - The revision state (asserted + derived edges with parents and weights)
"""
import numpy as np, hashlib, time
from pathlib import Path

def _seed(s): return int.from_bytes(hashlib.blake2b(s.encode(), digest_size=4).digest(), "little")
def phasor(n, D):
    rng = np.random.default_rng(_seed(n))
    return np.exp(1j*rng.uniform(-np.pi, np.pi, D).astype(np.float32)).astype(np.complex64)
def bind_dir(a, b):   return (a*np.roll(b, 1)).astype(np.complex64)
def unbind_dir(a, c): return np.roll(c*np.conj(a), -1).astype(np.complex64)
def sim(a, b):
    return float(np.real(np.vdot(a, b))/(np.linalg.norm(a)*np.linalg.norm(b)+1e-12))
def hdr(t): print(f"\n{'='*72}\n{t}\n{'='*72}")


class RevisionMemory:
    def __init__(self, D=1024):
        self.D = D
        self._items = {}
        self.asserted = {}   # (a, b) -> weight
        self.derived = {}    # (a, b) -> {"weight": float, "parents": [(a,b),(b,c)]}
        self.trace = np.zeros(D, dtype=np.complex64)

    def _node(self, name):
        v = self._items.get(name)
        if v is None:
            v = phasor(f"node::{name}", self.D)
            self._items[name] = v
        return v

    def _binding(self, a, b):
        return bind_dir(self._node(a), self._node(b))

    def _rebuild_trace(self):
        t = np.zeros(self.D, dtype=np.complex64)
        for (a, b), w in self.asserted.items():
            t = t + w * self._binding(a, b)
        for (a, b), e in self.derived.items():
            t = t + e["weight"] * self._binding(a, b)
        self.trace = t

    # ---------------- assertions ----------------
    def assert_edge(self, a, b, weight=1.0):
        self._node(a); self._node(b)
        self.asserted[(a, b)] = weight
        self._rebuild_trace()
        return self

    def forget(self, a, b, rederive=True):
        if (a, b) not in self.asserted:
            return {"ok": False}
        del self.asserted[(a, b)]
        self.derived = {}
        self._rebuild_trace()
        if rederive:
            self.derive()
        return {"ok": True, "asserted_remaining": len(self.asserted),
                "derived_after": len(self.derived)}

    # ---------------- derivation ----------------
    def derive(self, max_rounds=30):
        """Iterate forward chaining to fixpoint.  Weight = min(parent weights)."""
        for _ in range(max_rounds):
            all_edges = dict(self.asserted)
            for k, v in self.derived.items():
                all_edges[k] = v["weight"]
            new = {}
            for (a, b), w_ab in all_edges.items():
                for (b2, c), w_bc in all_edges.items():
                    if b2 != b or a == c: continue
                    if (a, c) in all_edges: continue
                    w = min(w_ab, w_bc)
                    if (a, c) not in new or new[(a, c)]["weight"] < w:
                        new[(a, c)] = {"weight": w, "parents": [(a, b), (b, c)]}
            if not new: break
            for k, v in new.items():
                if k not in self.derived or self.derived[k]["weight"] < v["weight"]:
                    self.derived[k] = v
        self._rebuild_trace()
        return len(self.derived)

    # ---------------- substrate queries ----------------
    def query(self, a, b):
        u = unbind_dir(self._node(a), self.trace)
        return sim(u, self._node(b))

    def _outgoing(self, a):
        return sorted({b for (x, b) in list(self.asserted) + list(self.derived) if x == a})

    def chain(self, start, n=10):
        path = [start]; cur = start; seen = {start}
        for _ in range(n):
            outs = self._outgoing(cur)
            if not outs: break
            scored = [(x, self.query(cur, x)) for x in outs]
            scored.sort(key=lambda t: -t[1])
            cur = scored[0][0]
            if cur in seen: break
            seen.add(cur); path.append(cur)
        return path

    def chain_quality(self, path):
        if len(path) < 2: return 1.0
        return min(self.query(path[i], path[i+1]) for i in range(len(path)-1))

    # ---------------- explainability ----------------
    def explain(self, a, b, depth=0):
        """Return a derivation tree."""
        if depth > 12:
            return {"edge": (a, b), "kind": "depth_limit"}
        if (a, b) in self.asserted:
            return {"edge": (a, b), "kind": "asserted",
                    "weight": self.asserted[(a, b)]}
        if (a, b) not in self.derived:
            return {"edge": (a, b), "kind": "unknown"}
        e = self.derived[(a, b)]
        (p1a, p1b), (p2a, p2b) = e["parents"]
        return {
            "edge": (a, b),
            "kind": "derived",
            "weight": e["weight"],
            "via": [self.explain(p1a, p1b, depth+1),
                    self.explain(p2a, p2b, depth+1)],
        }

    # ---------------- introspection ----------------
    def edges(self):
        return sorted(set(self.asserted) | set(self.derived))

    def stats(self):
        return {
            "nodes": len(self._items),
            "asserted": len(self.asserted),
            "derived": len(self.derived),
        }

    def __repr__(self):
        s = self.stats()
        return (f"RevisionMemory(D={self.D}, nodes={s['nodes']}, "
                f"asserted={s['asserted']}, derived={s['derived']})")


# ============================================================
def hdr(t): print(f"\n{'='*72}\n{t}\n{'='*72}")

def test_1_basic():
    hdr("T1. Basic derivation on a chain")
    mem = RevisionMemory(D=1024)
    for a, b in [("a","b"), ("b","c"), ("c","d"), ("d","e")]:
        mem.assert_edge(a, b)
    print(f"  Asserted: {sorted(mem.asserted.keys())}")
    mem.derive()
    print(f"  After derive: {sorted(mem.derived.keys())}")
    print(f"  Query a→c: sim={mem.query('a','c'):+.4f}")
    print(f"  Query a→e: sim={mem.query('a','e'):+.4f}")
    print(f"  Chain a: {' → '.join(mem.chain('a'))}")


def test_2_multipath():
    hdr("T2. Multi-path derivation — alternate route survives removal")
    mem = RevisionMemory(D=1024)
    for a, b in [("a","b"), ("b","c"), ("a","x"), ("x","c")]:
        mem.assert_edge(a, b)
    mem.derive()
    print(f"  Asserted: {sorted(mem.asserted.keys())}")
    print(f"  Derived:  {sorted(mem.derived.keys())}")
    print(f"  (a,c) weight: {mem.derived.get(('a','c'), {}).get('weight', '?')}")

    print(f"\n  forget a→b:")
    mem.forget("a", "b")
    print(f"  Asserted: {sorted(mem.asserted.keys())}")
    print(f"  Derived:  {sorted(mem.derived.keys())}")
    print(f"  (a,c) still derived? {('a','c') in mem.derived}")
    if ('a','c') in mem.derived:
        print(f"  (a,c) parents: {mem.derived[('a','c')]['parents']}")

    print(f"\n  forget a→x:")
    mem.forget("a", "x")
    print(f"  Asserted: {sorted(mem.asserted.keys())}")
    print(f"  Derived:  {sorted(mem.derived.keys())}")
    print(f"  (a,c) still derived? {('a','c') in mem.derived}")


def test_3_chain_quality():
    hdr("T3. Chain quality along walk")
    mem = RevisionMemory(D=2048)
    # chain of length 8
    nodes = list("abcdefghi")
    for i in range(len(nodes)-1):
        mem.assert_edge(nodes[i], nodes[i+1])
    mem.derive()
    path = mem.chain("a", n=10)
    print(f"  Path: {' → '.join(path)}")
    print(f"  Chain quality (min sim along walk): {mem.chain_quality(path):+.4f}")

    # now add noise by scattering many unrelated edges
    rng = np.random.default_rng(0)
    for _ in range(100):
        x, y = f"x{rng.integers(50)}", f"y{rng.integers(50)}"
        mem.assert_edge(x, y)
    mem.derive()
    path2 = mem.chain("a", n=10)
    print(f"  After 100 noise edges:")
    print(f"    Path: {' → '.join(path2)}")
    print(f"    Chain quality: {mem.chain_quality(path2):+.4f}")


def test_4_explain():
    hdr("T4. Explainability — why is (a,e)?")
    mem = RevisionMemory(D=1024)
    for a, b in [("a","b"), ("b","c"), ("c","d"), ("d","e")]:
        mem.assert_edge(a, b)
    mem.derive()
    tree = mem.explain("a", "e")
    def show(t, indent=0):
        pad = "  " * indent
        kind = t["kind"]
        e = t.get("edge")
        w = t.get("weight")
        print(f"{pad}{e} [{kind}]" + (f" w={w:.2f}" if w is not None else ""))
        if kind == "derived":
            for child in t["via"]:
                show(child, indent+1)
    show(tree)


def test_5_scaling():
    hdr("T5. Derivation cost vs chain length")
    print(f"  {'depth':>6}  {'asserted':>10}  {'derived':>9}  {'closure':>9}  "
          f"{'chain_q':>9}")
    for depth in [3, 5, 8, 12, 18]:
        mem = RevisionMemory(D=2048)
        nodes = [f"n{i}" for i in range(depth)]
        for i in range(depth-1):
            mem.assert_edge(nodes[i], nodes[i+1])
        mem.derive()
        closure_expected = depth * (depth-1) // 2
        path = mem.chain(nodes[0], n=depth+2)
        q = mem.chain_quality(path)
        print(f"  {depth:>6}  {len(mem.asserted):>10}  "
              f"{len(mem.derived):>9}  {closure_expected:>9}  {q:>+9.4f}")


def test_6_weighted():
    hdr("T6. Weighted assertions — weak links degrade derivation")
    mem = RevisionMemory(D=1024)
    mem.assert_edge("a", "b", weight=1.0)
    mem.assert_edge("b", "c", weight=0.5)
    mem.assert_edge("c", "d", weight=0.2)
    mem.derive()
    print(f"  Asserted weights: "
          f"a→b={mem.asserted[('a','b')]:.2f}  "
          f"b→c={mem.asserted[('b','c')]:.2f}  "
          f"c→d={mem.asserted[('c','d')]:.2f}")
    for edge in [("a","c"), ("b","d"), ("a","d")]:
        if edge in mem.derived:
            print(f"  Derived {edge}: weight={mem.derived[edge]['weight']:.2f}  "
                  f"parents={mem.derived[edge]['parents']}")


if __name__ == "__main__":
    t0 = time.time()
    test_1_basic()
    test_2_multipath()
    test_3_chain_quality()
    test_4_explain()
    test_5_scaling()
    test_6_weighted()
    print(f"\nRuntime: {time.time()-t0:.2f} s")