"""
holo-program
============
A memory where programs are first-class items, membranes are nested
regions, and writes can be cryptographically keyed.

Three layers, one substrate:

    Program layer      Operations are stored as key-value pairs:
                       key = random vector derived from (op, in)
                       value = bind(out, out_val)
                       Execution is unbind-then-classify. Chaining is
                       composition.

    Membrane layer     Regions have parents. Reads inherit from child
                       to parent. Writes isolate to the innermost
                       membrane. Dissolution merges a child into its
                       parent (P-systems operation). Movement transfers
                       an object between membranes.

    Cryptographic layer
                       Item vectors are derived from a key. Without
                       the key, a trace is noise. A commitment scheme
                       lets a store prove that an item was written
                       without revealing the item.

Everything is addressable through `ask()`.

Requires: numpy.
"""

from __future__ import annotations

import argparse
import json
import os
import zlib
from typing import Any, Dict, List, Optional, Tuple

import numpy as np


D_DEFAULT = 2048


# ============================================================
# Primitives
# ============================================================
def rvec(seed: int, d: int = D_DEFAULT) -> np.ndarray:
    r = np.random.default_rng(seed)
    return np.exp(1j * r.uniform(0, 2 * np.pi, d))


def povec(seed: int, d: int = D_DEFAULT) -> np.ndarray:
    v = rvec(seed, d)
    F = np.fft.fft(v)
    return np.fft.ifft(F / np.abs(F))


def bind(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    return np.fft.ifft(np.fft.fft(a) * np.fft.fft(b))


def unbind(a: np.ndarray, c: np.ndarray) -> np.ndarray:
    return np.fft.ifft(np.fft.fft(c) * np.conj(np.fft.fft(a)))


def norm(v: np.ndarray) -> np.ndarray:
    F = np.fft.fft(v)
    return np.fft.ifft(F / (np.abs(F) + 1e-12))


def cos(a: np.ndarray, b: np.ndarray) -> float:
    na = np.linalg.norm(a)
    nb = np.linalg.norm(b)
    if na < 1e-12 or nb < 1e-12:
        return 0.0
    return float(np.real(np.vdot(a, b)) / (na * nb))


def project(u: np.ndarray, v: np.ndarray) -> float:
    nv2 = float(np.real(np.vdot(v, v)))
    if nv2 < 1e-12:
        return 0.0
    return float(np.real(np.vdot(u, v)) / nv2)


def cos_batch(u: np.ndarray, X: np.ndarray) -> np.ndarray:
    X = X.reshape(len(X), -1)
    u = u.ravel()
    return np.real(X @ np.conj(u)) / (
        np.linalg.norm(X, axis=1) * np.linalg.norm(u) + 1e-12)


# ============================================================
# Cryptographic layer
# ============================================================
def keyed_seed(key: Optional[str], label: str) -> int:
    s = f"{key}::{label}" if key is not None else f"_::{label}"
    return zlib.crc32(s.encode()) & 0xffffffff


def make_vector(key: Optional[str], label: str, d: int) -> np.ndarray:
    return povec(keyed_seed(key, label), d)


# ============================================================
# Global role vectors (cached per dimension)
# ============================================================
_ROLES_CACHE: Dict[int, Dict[str, np.ndarray]] = {}


def get_roles(d: int) -> Dict[str, np.ndarray]:
    if d not in _ROLES_CACHE:
        _ROLES_CACHE[d] = {
            "op":   povec(1000, d),
            "in":   povec(1001, d),
            "out":  povec(1002, d),
            "salt": povec(1003, d),
        }
    return _ROLES_CACHE[d]


# ============================================================
# Membrane
# ============================================================
class Membrane:
    """
    A region of the substrate with an optional parent.
    """

    def __init__(self, name: str, d: int, key: Optional[str] = None,
                 parent: Optional["Membrane"] = None,
                 threshold: float = 0.05):
        self.name = name
        self.d = d
        self.key = key
        self.parent = parent
        self.threshold = threshold
        self.roles = get_roles(d)
        self.trace = np.zeros(d, dtype=complex)
        self.ops_trace = np.zeros(d, dtype=complex)
        self.weights: Dict[str, float] = {}
        self.labels: set = set()
        self.operations: List[Tuple[str, str, str]] = []

    def _vec(self, label: str) -> np.ndarray:
        return make_vector(self.key, label, self.d)

    def _op_key(self, op_name: str, in_val: str) -> np.ndarray:
        """
        Random key for an operation, derived from the pair.

        Fully random: no two distinct pairs share a component, so
        cross-talk between operations with the same op name is
        negligible. Same computation in add_op and apply.
        """
        pair_id = f"__op__{op_name}::{in_val}"
        return self._vec(pair_id)

    # ---------------- items ----------------
    def store(self, label: str, weight: float = 1.0) -> None:
        v = self._vec(label)
        self.trace += weight * bind(v, v)
        self.weights[label] = self.weights.get(label, 0.0) + weight
        self.labels.add(label)

    def similarity(self, label: str) -> float:
        v = self._vec(label)
        u = unbind(v, self.trace)
        return project(u, v)

    def query_inherited(self, label: str) -> float:
        sim = self.similarity(label)
        if sim >= self.threshold:
            return sim
        if self.parent is not None:
            return self.parent.query_inherited(label)
        return sim

    # ---------------- operations ----------------
    def add_op(self, op_name: str, in_val: str, out_val: str) -> None:
        key = self._op_key(op_name, in_val)
        value = bind(self.roles["out"], self._vec(out_val))
        self.ops_trace += bind(key, value)
        self.operations.append((op_name, in_val, out_val))

    def apply(self, op_name: str, in_val: str) -> Optional[str]:
        if not self.operations:
            return None
        key = self._op_key(op_name, in_val)
        u = unbind(key, self.ops_trace)
        out = unbind(self.roles["out"], u)
        candidates = sorted({o[2] for o in self.operations})
        stack = np.stack([self._vec(c) for c in candidates])
        sims = np.abs(cos_batch(out, stack))
        best = int(np.argmax(sims))
        if sims[best] < self.threshold:
            return None
        return candidates[best]

    def find_applicable(self, in_val: str
                        ) -> List[Tuple[str, str, float]]:
        hits: List[Tuple[str, str, float]] = []
        for op_name, in_v, out_v in self.operations:
            if in_v != in_val:
                continue
            key = self._op_key(op_name, in_val)
            u = unbind(key, self.ops_trace)
            out = unbind(self.roles["out"], u)
            sim = project(out, self._vec(out_v))
            if sim >= self.threshold:
                hits.append((op_name, out_v, float(sim)))
        hits.sort(key=lambda h: -h[2])
        return hits

    # ---------------- cryptographic ----------------
    def commit(self, label: str, salt: str) -> np.ndarray:
        v_label = self._vec(label)
        v_salt = self._vec(salt)
        return bind(v_label, v_salt)

    def verify(self, label: str, salt: str,
               commitment: np.ndarray) -> bool:
        expected = self.commit(label, salt)
        return cos(commitment, expected) > 0.99

    # ---------------- diagnostics ----------------
    def stats(self) -> Dict:
        return {
            "name": self.name,
            "d": self.d,
            "has_key": self.key is not None,
            "n_labels": len(self.labels),
            "n_operations": len(self.operations),
            "trace_magnitude": float(np.linalg.norm(self.trace)),
            "ops_magnitude": float(np.linalg.norm(self.ops_trace)),
        }


# ============================================================
# HoloProgram
# ============================================================
class HoloProgram:
    def __init__(self, d: int = D_DEFAULT,
                 key: Optional[str] = None,
                 threshold: float = 0.05):
        self.d = d
        self.key = key
        self.threshold = threshold
        self.membranes: Dict[str, Membrane] = {}
        self.add_membrane("root", parent=None)

    def add_membrane(self, name: str,
                     parent: Optional[str] = "root") -> None:
        if name in self.membranes:
            return
        parent_m = self.membranes.get(parent) if parent else None
        self.membranes[name] = Membrane(
            name=name, d=self.d, key=self.key,
            parent=parent_m, threshold=self.threshold,
        )

    def store(self, membrane: str, label: str,
              weight: float = 1.0) -> None:
        self.membranes[membrane].store(label, weight=weight)

    def add_op(self, membrane: str, op_name: str,
               in_val: str, out_val: str) -> None:
        self.membranes[membrane].add_op(op_name, in_val, out_val)

    def query(self, membrane: str, label: str) -> float:
        return self.membranes[membrane].query_inherited(label)

    def apply(self, membrane: str, op_name: str,
              in_val: str) -> Optional[str]:
        return self.membranes[membrane].apply(op_name, in_val)

    def run_chain(self, membrane: str, ops: List[str],
                  start: str) -> Dict:
        cur = start
        path = [cur]
        for op in ops:
            nxt = self.apply(membrane, op, cur)
            if nxt is None:
                return {"path": path, "ok": False, "failed_at": op}
            path.append(nxt)
            cur = nxt
        return {"path": path, "ok": True}

    def run_data_driven(self, membrane: str, start: str,
                        max_steps: int = 20) -> Dict:
        cur = start
        path = [cur]
        for _ in range(max_steps):
            hits = self.membranes[membrane].find_applicable(cur)
            if not hits:
                break
            op_name, out_val, sim = hits[0]
            path.append(out_val)
            cur = out_val
        return {"path": path, "steps": len(path) - 1}

    def commit(self, membrane: str, label: str,
               salt: str) -> np.ndarray:
        return self.membranes[membrane].commit(label, salt)

    def verify(self, membrane: str, label: str, salt: str,
               commitment: np.ndarray) -> bool:
        return self.membranes[membrane].verify(label, salt, commitment)

    def move(self, label: str, src: str, dst: str) -> bool:
        s = self.membranes[src]
        d = self.membranes[dst]
        if label not in s.labels:
            return False
        w = s.weights.get(label, 1.0)
        if w <= 0:
            return False
        v = s._vec(label)
        s.trace -= w * bind(v, v)
        s.weights[label] = 0.0
        s.labels.discard(label)
        d.store(label, weight=w)
        return True

    def dissolve(self, child: str, into: str) -> None:
        c = self.membranes[child]
        t = self.membranes[into]
        t.trace = t.trace + c.trace
        t.ops_trace = t.ops_trace + c.ops_trace
        for lbl, w in c.weights.items():
            t.weights[lbl] = t.weights.get(lbl, 0.0) + w
            t.labels.add(lbl)
        t.operations.extend(c.operations)
        for m in self.membranes.values():
            if m.parent is c:
                m.parent = t
        del self.membranes[child]

    def ask(self, query: str, top_k: int = 8) -> List[Dict]:
        hits: List[Dict] = []

        for m_name, mem in self.membranes.items():
            if query in mem.labels:
                sim = mem.query_inherited(query)
                hits.append({
                    "type": "item",
                    "membrane": m_name,
                    "label": query,
                    "similarity": sim,
                })

        for m_name, mem in self.membranes.items():
            for op_name, in_v, out_v in mem.operations:
                if query in (op_name, in_v, out_v):
                    role = ("op" if query == op_name else
                            "in" if query == in_v else "out")
                    hits.append({
                        "type": "operation",
                        "membrane": m_name,
                        "role": role,
                        "op": op_name,
                        "in": in_v,
                        "out": out_v,
                    })

        if query in self.membranes:
            hits.append({
                "type": "membrane",
                "membrane": query,
            })

        def sort_key(h):
            if h["type"] == "item":
                return (0, -h["similarity"])
            return (1, 0)

        hits.sort(key=sort_key)
        return hits[:top_k]

    def stats(self) -> Dict:
        return {
            "d": self.d,
            "has_key": self.key is not None,
            "threshold": self.threshold,
            "n_membranes": len(self.membranes),
            "membranes": {n: m.stats()
                          for n, m in self.membranes.items()},
        }

    def save_json(self, path: str) -> None:
        out = self.stats()
        out["membrane_detail"] = {
            n: {
                "labels": sorted(m.labels),
                "operations": [
                    {"op": o, "in": i, "out": r}
                    for o, i, r in m.operations
                ],
                "parent": m.parent.name if m.parent else None,
            }
            for n, m in self.membranes.items()
        }
        with open(path, "w") as f:
            json.dump(out, f, indent=2)
        print(f"  saved: {path}")


# ============================================================
# Helpers
# ============================================================
def section(title: str) -> None:
    print()
    print("=" * 78)
    print(title)
    print("=" * 78)


def self_test() -> bool:
    print("  self-test:")
    d = 2048
    a = povec(1, d)
    b = povec(2, d)
    ok1 = cos(unbind(a, bind(a, b)), b) > 0.99
    print(f"    bind/unbind identity:      {'PASS' if ok1 else 'FAIL'}")

    v1 = make_vector("key1", "label", d)
    v2 = make_vector("key2", "label", d)
    ok2 = cos(v1, v2) < 0.1
    print(f"    keys separate vectors:     {'PASS' if ok2 else 'FAIL'}")

    bb = 0.7 * bind(a, a)
    recovered = project(unbind(a, bb), a)
    ok3 = abs(recovered - 0.7) < 0.02
    print(f"    raw projection:            {'PASS' if ok3 else 'FAIL'}")

    # Fully random key-value storage: many pairs sharing nothing
    r = get_roles(d)
    trace = np.zeros(d, dtype=complex)
    pairs = []
    for i in range(50):
        key = norm(bind(r["op"], povec(2000 + i, d)))
        val = povec(3000 + i, d)
        trace += bind(key, val)
        pairs.append((key, val))
    correct = 0
    for key, val in pairs:
        rec = unbind(key, trace)
        if cos(rec, val) > cos(rec, pairs[0][1]):
            correct += 1
    ok4 = correct >= 48
    print(f"    random-key key-value:      "
          f"{'PASS' if ok4 else 'FAIL'}  ({correct}/50)")

    return ok1 and ok2 and ok3 and ok4


# ============================================================
# DEMO 1. Program execution (chains)
# ============================================================
def demo_program_chains():
    section("DEMO 1. PROGRAM EXECUTION — chained operations")
    prog = HoloProgram(d=2048)

    for i in range(20):
        prog.add_op("root", "inc", str(i), str(i + 1))
    for i in range(20):
        prog.add_op("root", "double", str(i), str(2 * i))
    for i in range(10):
        prog.add_op("root", "square", str(i), str(i * i))

    print("  stored: 20 inc, 20 double, 10 square operations")
    print()

    chains = [
        (["inc", "inc", "double"], "3"),
        (["square", "inc"], "4"),
        (["double", "inc", "inc"], "5"),
        (["inc", "square"], "2"),
    ]
    for ops, start in chains:
        result = prog.run_chain("root", ops, start)
        path_str = " -> ".join(result["path"])
        ok = "OK" if result["ok"] else f"FAILED at {result['failed_at']}"
        print(f"  [{', '.join(ops)}] from '{start}': {path_str}  {ok}")

    print()
    print("  single apply:")
    for op_name, in_val in [("inc", "7"), ("double", "7"), ("square", "7")]:
        out = prog.apply("root", op_name, in_val)
        print(f"    {op_name}({in_val}) -> {out}")


# ============================================================
# DEMO 2. Data-driven execution
# ============================================================
def demo_data_driven():
    section("DEMO 2. DATA-DRIVEN EXECUTION — apply any applicable op")
    prog = HoloProgram(d=2048)
    transitions = [
        ("t1", "cold", "warm"),
        ("t2", "warm", "hot"),
        ("t3", "hot", "boiling"),
        ("t4", "boiling", "evaporated"),
    ]
    for op, src, dst in transitions:
        prog.add_op("root", op, src, dst)

    result = prog.run_data_driven("root", "cold", max_steps=10)
    print(f"  start at 'cold', walk using whichever op matches:")
    print(f"    {' -> '.join(result['path'])}")
    print(f"    steps: {result['steps']}")

    print()
    result = prog.run_data_driven("root", "evaporated", max_steps=5)
    print(f"  start at 'evaporated' (no matching op):")
    print(f"    {' -> '.join(result['path'])}")
    print(f"    steps: {result['steps']}")


# ============================================================
# DEMO 3. Membranes with inheritance
# ============================================================
def demo_membranes():
    section("DEMO 3. MEMBRANES — nested regions with inheritance")
    prog = HoloProgram(d=2048)
    prog.add_membrane("alice", parent="root")
    prog.add_membrane("alice_kitchen", parent="alice")

    prog.store("root", "earth_round")
    prog.store("root", "water_wet")
    prog.store("alice", "alice_lives_in_paris")
    prog.store("alice", "alice_works_at_cafe")
    prog.store("alice_kitchen", "kitchen_floor_5")
    prog.store("alice_kitchen", "manager_is_bob")

    print("  membranes: root <- alice <- alice_kitchen")
    print()
    print(f"  {'query':>24}  {'root':>8}  {'alice':>8}  {'kitchen':>8}")
    for q in ["earth_round", "alice_lives_in_paris",
              "kitchen_floor_5", "manager_is_bob"]:
        r_root = prog.query("root", q)
        r_alice = prog.query("alice", q)
        r_kit = prog.query("alice_kitchen", q)
        print(f"  {q:>24}  {r_root:>+8.3f}  {r_alice:>+8.3f}  {r_kit:>+8.3f}")


# ============================================================
# DEMO 4. Communication and dissolution
# ============================================================
def demo_communication():
    section("DEMO 4. COMMUNICATION AND DISSOLUTION")
    prog = HoloProgram(d=2048)
    prog.add_membrane("inner", parent="root")
    prog.store("root", "common_knowledge")
    prog.store("inner", "private_note")

    print("  before move:")
    print(f"    root sees private_note:    {prog.query('root', 'private_note'):+.4f}")
    print(f"    inner sees private_note:   {prog.query('inner', 'private_note'):+.4f}")

    prog.move("private_note", "inner", "root")
    print()
    print("  after move (inner -> root):")
    print(f"    root sees private_note:    {prog.query('root', 'private_note'):+.4f}")
    print(f"    inner sees private_note:   {prog.query('inner', 'private_note'):+.4f}")

    prog.add_membrane("middle", parent="root")
    prog.store("middle", "middle_fact")
    print()
    print("  before dissolution:")
    print(f"    root sees middle_fact:     {prog.query('root', 'middle_fact'):+.4f}")
    print(f"    middle sees middle_fact:   {prog.query('middle', 'middle_fact'):+.4f}")

    prog.dissolve("middle", "root")
    print()
    print("  after dissolution (middle -> root):")
    print(f"    root sees middle_fact:     {prog.query('root', 'middle_fact'):+.4f}")
    print(f"    membranes remaining: {list(prog.membranes.keys())}")


# ============================================================
# DEMO 5. Cryptographic layer
# ============================================================
def demo_crypto():
    section("DEMO 5. CRYPTOGRAPHIC LAYER — keyed codebook")
    print("  Store items with key 'secret_A'. Copy the trace to two")
    print("  other programs: one with the same key, one with a wrong")
    print("  key. Query both.")
    print()

    p_correct = HoloProgram(d=2048, key="secret_A")
    p_correct.store("root", "classified_meeting_place")
    p_correct.store("root", "classified_meeting_time")

    p_wrong = HoloProgram(d=2048, key="wrong_key")
    p_wrong.membranes["root"].trace = \
        p_correct.membranes["root"].trace.copy()

    candidates = ["classified_meeting_place", "classified_meeting_time",
                  "paris", "3pm", "unrelated_word"]

    print(f"  {'candidate':>30}  {'correct key':>14}  {'wrong key':>12}")
    for c in candidates:
        sim_c = p_correct.query("root", c)
        sim_w = p_wrong.query("root", c)
        print(f"  {c:>30}  {sim_c:>+14.4f}  {sim_w:>+12.4f}")


# ============================================================
# DEMO 6. Commitment scheme
# ============================================================
def demo_commit():
    section("DEMO 6. COMMITMENT SCHEME")
    prog = HoloProgram(d=2048, key="K")
    prog.store("root", "secret_value")

    commitment = prog.commit("root", "secret_value", salt="random_salt_123")
    print(f"  commitment vector shape:     {commitment.shape}")
    print(f"  commitment magnitude:        {np.linalg.norm(commitment):.4f}")
    print()
    print("  verification:")
    ok1 = prog.verify("root", "secret_value", "random_salt_123", commitment)
    ok2 = prog.verify("root", "secret_value", "wrong_salt", commitment)
    ok3 = prog.verify("root", "different_value", "random_salt_123", commitment)
    print(f"    correct (label, salt):  {ok1}")
    print(f"    wrong salt:             {ok2}")
    print(f"    wrong label:            {ok3}")


# ============================================================
# DEMO 7. Unified ask()
# ============================================================
def demo_ask():
    section("DEMO 7. UNIFIED ask() — non-flat interface")
    prog = HoloProgram(d=2048)
    prog.add_membrane("alice", parent="root")
    prog.add_membrane("alice_kitchen", parent="alice")

    prog.store("root", "earth_round")
    prog.store("root", "water_wet")
    prog.store("alice", "alice_lives_in_paris")
    prog.store("alice", "alice_works_at_cafe")
    prog.store("alice_kitchen", "kitchen_floor_5")
    prog.store("alice_kitchen", "manager_is_bob")

    prog.add_op("root", "serve", "customer", "coffee")
    prog.add_op("root", "greet", "customer", "hello")
    prog.add_op("alice_kitchen", "brew", "coffee", "ready")

    for query in ["alice", "alice_lives_in_paris", "coffee",
                  "alice_kitchen", "manager_is_bob"]:
        hits = prog.ask(query, top_k=6)
        print(f"  ask('{query}'):")
        if not hits:
            print(f"    (no hits)")
        for h in hits:
            if h["type"] == "item":
                print(f"    [item]      membrane={h['membrane']:<14} "
                      f"sim={h['similarity']:+.3f}")
            elif h["type"] == "operation":
                print(f"    [operation] membrane={h['membrane']:<14} "
                      f"role={h['role']:<4} "
                      f"op={h['op']}({h['in']}) -> {h['out']}")
            elif h["type"] == "membrane":
                print(f"    [membrane]  {h['membrane']}")
        print()


# ============================================================
# DEMO 8. Programs as items
# ============================================================
def demo_program_as_item():
    section("DEMO 8. PROGRAMS AS ITEMS")
    prog = HoloProgram(d=2048)
    prog.add_op("root", "inc", "1", "2")
    prog.add_op("root", "inc", "2", "3")
    prog.add_op("root", "double", "3", "6")

    print("  The substrate holds operations alongside items.")
    print("  Query for 'inc' through ask():")
    hits = prog.ask("inc", top_k=10)
    for h in hits:
        if h["type"] == "operation":
            print(f"    {h['op']}({h['in']}) -> {h['out']}  "
                  f"(membrane={h['membrane']})")

    print()
    print("  Query for '3' (input and output of operations):")
    hits = prog.ask("3", top_k=10)
    for h in hits:
        if h["type"] == "operation":
            print(f"    {h['op']}({h['in']}) -> {h['out']}  "
                  f"role={h['role']}")


# ============================================================
# Main
# ============================================================
def run_all_demos(output_dir: str = "holo_program_out") -> None:
    os.makedirs(output_dir, exist_ok=True)
    print("=" * 78)
    print("HOLO-PROGRAM")
    print("Programs as items, membranes as regions, keys as gates.")
    print("=" * 78)

    demos = [
        demo_program_chains,
        demo_data_driven,
        demo_membranes,
        demo_communication,
        demo_crypto,
        demo_commit,
        demo_ask,
        demo_program_as_item,
    ]
    for d in demos:
        try:
            d()
        except Exception as e:
            print(f"  EXCEPTION in {d.__name__}: {e}")

    section("OUTPUT")
    prog = HoloProgram(d=2048, key="demo_key")
    prog.add_membrane("alice", parent="root")
    prog.store("root", "earth_round")
    prog.store("alice", "alice_lives_in_paris")
    prog.add_op("root", "inc", "1", "2")
    prog.save_json(os.path.join(output_dir, "state.json"))

    print()
    print("=" * 78)
    print(f"All outputs saved to: {output_dir}/")
    print("=" * 78)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Programs as items, membranes as regions, "
                    "keys as gates.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Library usage:

    from holo_program import HoloProgram

    prog = HoloProgram(d=2048, key=None)
    prog.add_membrane("alice", parent="root")
    prog.store("root", "apple")
    prog.add_op("root", "inc", "0", "1")
    prog.run_chain("root", ["inc", "inc"], "0")
        """)
    parser.add_argument("--output", default="holo_program_out",
                        help="output directory (default: holo_program_out)")
    args = parser.parse_args()

    print("=" * 78)
    print("SELF-TEST")
    print("=" * 78)
    if not self_test():
        print("  primitive test failed; continuing anyway")

    run_all_demos(output_dir=args.output)


if __name__ == "__main__":
    main()