"""
holo-four-axis
==============
A single memory store parameterized on four orthogonal axes.

    Axis 1: dimension (d)             capacity
    Axis 2: algebra (kind)            addressability, address shape
    Axis 3: threshold                 readout confidence
    Axis 4: consolidation_rate        forgetting behavior over time

Every axis is measured. None of them is fixed by another. The user
chooses a profile and gets a store whose behavior follows from the
profile.

    - Small d, tight threshold, fast consolidation:
        a fast-shallow working memory

    - Large d, loose threshold, slow consolidation:
        a wide-persistent long-term memory

    - Directed algebra, tight threshold, medium consolidation:
        a sequence memory with confidence

    - Anti algebra, loose threshold, slow consolidation:
        a contradiction-aware knowledge base

The four axes are the substrate's design space, extracted from the
prior experiments in the series.

Usage
-----

    from holo_four_axis import FourAxisMemory

    mem = FourAxisMemory(
        d=2048,
        kind="item",
        threshold=0.05,
        consolidation_rate=0.5,
    )
    mem.add_region("items", kind="item")
    mem.add_region("events", kind="chain")
    mem.add_region("claims", kind="anti")

    mem.store("items", "apple")
    mem.store("items", "orange")
    mem.chain("events", ["wake", "coffee", "work", "sleep"])
    mem.store("claims", "earth_round")
    mem.deny("claims", "earth_flat")

    print(mem.query("items", "apple"))
    print(mem.query("claims", "earth_flat"))

CLI
---

    python holo_four_axis.py
    python holo_four_axis.py --output results/

Requires: numpy. matplotlib optional (only for the summary plot).
"""

from __future__ import annotations

import argparse
import json
import os
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Optional, Tuple, Union

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 bind_dir(a: np.ndarray, b: np.ndarray, shift: int = 1) -> np.ndarray:
    return np.fft.ifft(np.fft.fft(a) * np.roll(np.fft.fft(b), shift))


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 unbind_dir(a: np.ndarray, c: np.ndarray, shift: int = 1) -> np.ndarray:
    return np.fft.ifft(np.roll(np.fft.fft(c) * np.conj(np.fft.fft(a)), -shift))


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:
    """Raw projection of u onto v, normalized by ||v||^2."""
    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)


# ============================================================
# Region
# ============================================================
@dataclass
class Region:
    name: str
    kind: str  # "item", "anti", "chain"
    d: int
    seed: int
    trace_yes: np.ndarray
    trace_no: Optional[np.ndarray] = None
    trace_prev: Optional[np.ndarray] = None
    items: Dict[str, np.ndarray] = field(default_factory=dict)
    weight_yes: Dict[str, float] = field(default_factory=dict)
    weight_no: Dict[str, float] = field(default_factory=dict)
    edges: List[Tuple[int, int]] = field(default_factory=list)
    label_index: List[str] = field(default_factory=list)
    counter: int = 1

    def get_or_create_vector(self, label: str) -> np.ndarray:
        if label not in self.items:
            v = povec(self.seed + self.counter, self.d)
            self.counter += 1
            self.items[label] = v
            self.weight_yes.setdefault(label, 0.0)
            self.weight_no.setdefault(label, 0.0)
        return self.items[label]


# ============================================================
# FourAxisMemory
# ============================================================
class FourAxisMemory:
    """
    A memory store with four parameterized axes.

    Regions are independent traces. Each region has its own kind
    (item, anti, chain). The store routes operations to the region
    and applies the readout threshold to every query. Consolidation
    runs over all regions and decays non-anchor weights.
    """

    def __init__(self, d: int = D_DEFAULT, kind: str = "item",
                 threshold: float = 0.05,
                 consolidation_rate: float = 0.0,
                 anchor_quantile: float = 0.8,
                 seed: int = 0):
        # Axis 1: dimension
        self.d = d
        # Axis 2: default algebra for the default region
        self.default_kind = kind
        # Axis 3: readout threshold
        self.threshold = threshold
        # Axis 4: consolidation rate
        self.consolidation_rate = consolidation_rate
        self.anchor_quantile = anchor_quantile
        self.seed = seed
        self.regions: Dict[str, Region] = {}
        self._seed_counter = 1
        self.consolidation_cycles = 0

        # Create a default region if kind is provided
        self.add_region("default", kind=kind)

    # --------------------------------------------------------
    # Region management
    # --------------------------------------------------------
    def add_region(self, name: str, kind: str = "item") -> None:
        if name in self.regions:
            return
        trace_yes = np.zeros(self.d, dtype=complex)
        trace_no = (np.zeros(self.d, dtype=complex)
                    if kind == "anti" else None)
        trace_prev = (np.zeros(self.d, dtype=complex)
                      if kind == "chain" else None)
        seed = self.seed + self._seed_counter * 1000
        self._seed_counter += 1
        self.regions[name] = Region(
            name=name, kind=kind, d=self.d, seed=seed,
            trace_yes=trace_yes, trace_no=trace_no,
            trace_prev=trace_prev,
        )

    def _region(self, name: str) -> Region:
        if name not in self.regions:
            raise KeyError(f"region '{name}' does not exist")
        return self.regions[name]

    # --------------------------------------------------------
    # Store / deny / chain
    # --------------------------------------------------------
    def store(self, region: str, label: str, weight: float = 1.0) -> None:
        r = self._region(region)
        v = r.get_or_create_vector(label)
        if r.kind == "item":
            r.trace_yes += weight * bind(v, v)
            r.weight_yes[label] = r.weight_yes.get(label, 0.0) + weight
        elif r.kind == "anti":
            r.trace_yes += weight * bind(v, v)
            r.weight_yes[label] = r.weight_yes.get(label, 0.0) + weight
        elif r.kind == "chain":
            # In a chain region, "store" adds a node without an edge.
            r.weight_yes[label] = r.weight_yes.get(label, 0.0) + weight

    def deny(self, region: str, label: str, weight: float = 1.0) -> None:
        r = self._region(region)
        if r.kind != "anti":
            raise ValueError(
                f"region '{region}' has kind '{r.kind}'; "
                f"deny is only valid for anti regions")
        v = r.get_or_create_vector(label)
        r.trace_no += weight * bind(v, v)
        r.weight_no[label] = r.weight_no.get(label, 0.0) + weight

    def store_many(self, region: str, labels: List[str],
                   weight: float = 1.0) -> None:
        for lb in labels:
            self.store(region, lb, weight=weight)

    def chain(self, region: str, labels: List[str],
              weight: float = 1.0) -> None:
        r = self._region(region)
        if r.kind != "chain":
            raise ValueError(
                f"region '{region}' has kind '{r.kind}'; "
                f"chain is only valid for chain regions")
        indices = []
        for lb in labels:
            v = r.get_or_create_vector(lb)
            idx = len(r.label_index)
            r.label_index.append(lb)
            r.weight_yes[lb] = r.weight_yes.get(lb, 0.0) + weight
            indices.append(idx)
        for a, b in zip(indices[:-1], indices[1:]):
            va = r.items[r.label_index[a]]
            vb = r.items[r.label_index[b]]
            r.trace_yes += weight * bind_dir(va, vb, shift=1)
            r.trace_prev += weight * bind_dir(vb, va, shift=1)
            r.edges.append((a, b))

    # --------------------------------------------------------
    # Query
    # --------------------------------------------------------
    def query(self, region: str, label: str) -> Dict:
        r = self._region(region)
        if label not in r.items:
            return {
                "region": region, "label": label,
                "verdict": "UNKNOWN", "similarity": 0.0,
                "confidence": 0.0, "raw_yes": 0.0, "raw_no": 0.0,
            }
        v = r.items[label]

        if r.kind == "item":
            u = unbind(v, r.trace_yes)
            sim = project(u, v)
            if sim >= self.threshold:
                return {
                    "region": region, "label": label,
                    "verdict": "YES", "similarity": sim,
                    "confidence": sim,
                    "raw_yes": sim, "raw_no": 0.0,
                }
            else:
                return {
                    "region": region, "label": label,
                    "verdict": "UNKNOWN", "similarity": sim,
                    "confidence": sim,
                    "raw_yes": sim, "raw_no": 0.0,
                }

        elif r.kind == "anti":
            u_y = unbind(v, r.trace_yes)
            u_n = unbind(v, r.trace_no)
            ry = project(u_y, v)
            rn = project(u_n, v)
            wy = r.weight_yes.get(label, 0.0)
            wn = r.weight_no.get(label, 0.0)
            yes_active = wy > 0 and ry >= self.threshold
            no_active = wn > 0 and rn >= self.threshold
            if yes_active and no_active:
                verdict = "AMBIGUOUS"
            elif yes_active:
                verdict = "YES"
            elif no_active:
                verdict = "NO"
            else:
                verdict = "UNKNOWN"
            return {
                "region": region, "label": label,
                "verdict": verdict, "similarity": max(ry, rn),
                "confidence": max(ry, rn),
                "raw_yes": ry, "raw_no": rn,
            }

        elif r.kind == "chain":
            # Query a chain region returns whether the label is a known
            # node. Retrieval of successors and predecessors is a
            # separate operation.
            if label in r.items:
                return {
                    "region": region, "label": label,
                    "verdict": "YES", "similarity": 1.0,
                    "confidence": 1.0,
                    "raw_yes": 1.0, "raw_no": 0.0,
                }

        return {
            "region": region, "label": label,
            "verdict": "UNKNOWN", "similarity": 0.0,
            "confidence": 0.0, "raw_yes": 0.0, "raw_no": 0.0,
        }

    def query_next(self, region: str, label: str) -> Optional[str]:
        r = self._region(region)
        if r.kind != "chain":
            raise ValueError(f"query_next requires chain region, "
                             f"got '{r.kind}'")
        if label not in r.items:
            return None
        v = r.items[label]
        u = unbind_dir(v, r.trace_yes, shift=1)
        labels = list(r.items.keys())
        if not labels:
            return None
        stack = np.stack([r.items[lb] for lb in labels])
        sims = np.abs(cos_batch(u, stack))
        sims[labels.index(label)] = -1.0
        best_idx = int(np.argmax(sims))
        if sims[best_idx] < self.threshold:
            return None
        return labels[best_idx]

    def query_prev(self, region: str, label: str) -> Optional[str]:
        r = self._region(region)
        if r.kind != "chain":
            raise ValueError(f"query_prev requires chain region, "
                             f"got '{r.kind}'")
        if label not in r.items:
            return None
        v = r.items[label]
        u = unbind_dir(v, r.trace_prev, shift=1)
        labels = list(r.items.keys())
        if not labels:
            return None
        stack = np.stack([r.items[lb] for lb in labels])
        sims = np.abs(cos_batch(u, stack))
        sims[labels.index(label)] = -1.0
        best_idx = int(np.argmax(sims))
        if sims[best_idx] < self.threshold:
            return None
        return labels[best_idx]

    def walk(self, region: str, start: str,
             max_steps: int = 20) -> List[str]:
        chain = [start]
        visited = {start}
        cur = start
        for _ in range(max_steps):
            nxt = self.query_next(region, cur)
            if nxt is None or nxt in visited:
                break
            chain.append(nxt)
            visited.add(nxt)
            cur = nxt
        return chain

    # --------------------------------------------------------
    # Consolidation
    # --------------------------------------------------------
    def consolidate(self, cycles: int = 1) -> Dict:
        """
        Run consolidation. Non-anchor weights decay by
        consolidation_rate each cycle; anchors (weights above the
        anchor quantile) do not decay.
        """
        if self.consolidation_rate <= 0:
            return {"cycles": 0, "decayed": 0, "anchors_preserved": 0}

        total_decayed = 0
        total_anchors = 0
        for _ in range(cycles):
            for r in self.regions.values():
                weights = list(r.weight_yes.values())
                if not weights:
                    continue
                q = np.quantile(weights, self.anchor_quantile)
                for label in list(r.weight_yes.keys()):
                    if r.weight_yes[label] >= q:
                        total_anchors += 1
                        continue
                    # Decay this item's contribution to the trace
                    decay = self.consolidation_rate
                    if r.weight_yes[label] > 0 and label in r.items:
                        v = r.items[label]
                        r.trace_yes -= (decay * r.weight_yes[label]
                                        * bind(v, v))
                        r.weight_yes[label] *= (1 - decay)
                        total_decayed += 1
                    if (r.kind == "anti" and r.weight_no.get(label, 0) > 0
                            and label in r.items):
                        v = r.items[label]
                        r.trace_no -= (decay * r.weight_no[label]
                                       * bind(v, v))
                        r.weight_no[label] *= (1 - decay)
            self.consolidation_cycles += 1

        return {
            "cycles": cycles,
            "decayed": total_decayed,
            "anchors_preserved": total_anchors,
        }

    # --------------------------------------------------------
    # Diagnostics
    # --------------------------------------------------------
    def stats(self) -> Dict:
        out = {
            "d": self.d,
            "default_kind": self.default_kind,
            "threshold": self.threshold,
            "consolidation_rate": self.consolidation_rate,
            "consolidation_cycles": self.consolidation_cycles,
            "regions": {},
        }
        for name, r in self.regions.items():
            n_items = len(r.items)
            n_pos = sum(1 for w in r.weight_yes.values() if w > 0)
            n_neg = sum(1 for w in r.weight_no.values() if w > 0)
            out["regions"][name] = {
                "kind": r.kind,
                "n_items": n_items,
                "n_positive": n_pos,
                "n_negative": n_neg,
                "n_edges": len(r.edges),
                "magnitude_yes": float(np.linalg.norm(r.trace_yes)),
                "magnitude_no": (
                    float(np.linalg.norm(r.trace_no))
                    if r.trace_no is not None else 0.0),
            }
        return out

    def save_json(self, path: str) -> None:
        out = self.stats()
        out["regions_detail"] = {
            name: {
                "items": list(r.items.keys()),
                "weight_yes": {k: float(v)
                               for k, v in r.weight_yes.items()},
                "weight_no": {k: float(v)
                              for k, v in r.weight_no.items()},
                "edges": [[int(a), int(b)] for a, b in r.edges],
                "label_index": list(r.label_index),
            }
            for name, r in self.regions.items()
        }
        with open(path, "w") as f:
            json.dump(out, f, indent=2)
        print(f"  saved: {path}")

    def plot(self, path: str) -> None:
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
        except ImportError:
            print("  matplotlib not available; skipping plot")
            return
        region_names = list(self.regions.keys())
        if not region_names:
            return
        fig, axes = plt.subplots(1, len(region_names),
                                 figsize=(4 * len(region_names), 3.5),
                                 squeeze=False)
        for ax, name in zip(axes[0], region_names):
            r = self.regions[name]
            labels = list(r.items.keys())
            if not labels:
                ax.text(0.5, 0.5, "(empty)", ha="center", va="center")
                ax.set_title(name)
                continue
            wy = [r.weight_yes.get(lb, 0.0) for lb in labels]
            wn = [r.weight_no.get(lb, 0.0) for lb in labels]
            x = np.arange(len(labels))
            ax.bar(x - 0.2, wy, 0.4, label="yes",
                   color="#2a9d3f", edgecolor="black", lw=0.4)
            ax.bar(x + 0.2, wn, 0.4, label="no",
                   color="#c44e4e", edgecolor="black", lw=0.4)
            ax.set_xticks(x)
            ax.set_xticklabels([lb[:10] for lb in labels],
                               rotation=45, ha="right", fontsize=6)
            ax.set_title(f"{name} ({r.kind})")
            ax.legend(fontsize=7)
            ax.grid(alpha=0.3, axis="y")
        plt.tight_layout()
        plt.savefig(path, dpi=110, bbox_inches="tight")
        plt.close()
        print(f"  plot: {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)
    bb = bind(a, a)
    ok1 = cos(unbind(a, bb), a) > 0.99
    print(f"    bind/unbind identity:      {'PASS' if ok1 else 'FAIL'}")
    bd = bind_dir(a, b, shift=1)
    ok2 = cos(unbind_dir(a, bd, shift=1), b) > 0.99
    print(f"    directed bind identity:    {'PASS' if ok2 else 'FAIL'}")
    ok3 = abs(project(unbind(a, 0.7 * bind(a, a)), a) - 0.7) < 0.02
    print(f"    raw projection:            {'PASS' if ok3 else 'FAIL'}")
    return ok1 and ok2 and ok3


# ============================================================
# AXIS 1. Dimension
# ============================================================
def axis1_dimension():
    section("AXIS 1. DIMENSION (capacity)")
    print(f"  {'D':>6}  {'K':>5}  {'hits':>8}")
    for d in [512, 1024, 2048]:
        for K in [20, 50, 100, 200]:
            mem = FourAxisMemory(d=d)
            for i in range(K):
                mem.store("default", f"item_{i}")
            hits = 0
            for i in range(K):
                r = mem.query("default", f"item_{i}")
                if r["verdict"] == "YES":
                    hits += 1
            print(f"  {d:>6}  {K:>5}  {hits:>4}/{K:<3}")


# ============================================================
# AXIS 2. Algebra
# ============================================================
def axis2_algebra():
    section("AXIS 2. ALGEBRA (address shape)")
    mem = FourAxisMemory(d=2048)
    mem.add_region("items", kind="item")
    mem.add_region("claims", kind="anti")
    mem.add_region("events", kind="chain")

    # Item region
    mem.store("items", "apple")
    mem.store("items", "orange")
    mem.store("items", "banana")

    # Anti region
    mem.store("claims", "earth_round")
    mem.deny("claims", "earth_flat")
    mem.store("claims", "water_is_wet")
    mem.deny("claims", "water_is_wet")

    # Chain region
    mem.chain("events", ["wake", "coffee", "work", "lunch",
                         "afternoon", "dinner", "sleep"])

    print("  item region:")
    for lb in ["apple", "orange", "banana", "pear"]:
        r = mem.query("items", lb)
        print(f"    {lb:>10}  {r['verdict']:>10}  "
              f"sim={r['similarity']:>+7.4f}")

    print()
    print("  anti region:")
    for lb in ["earth_round", "earth_flat", "water_is_wet",
               "sun_orbits_earth"]:
        r = mem.query("claims", lb)
        print(f"    {lb:>18}  {r['verdict']:>10}  "
              f"raw_yes={r['raw_yes']:>+6.3f}  "
              f"raw_no={r['raw_no']:>+6.3f}")

    print()
    print("  chain region:")
    chain = mem.walk("events", "wake")
    print(f"    {' -> '.join(chain)}")

    print()
    print("  same primitives, three algebras, three address shapes.")


# ============================================================
# AXIS 3. Threshold
# ============================================================
def axis3_threshold():
    section("AXIS 3. THRESHOLD (confidence gate)")
    print("  Store 30 items. Query with stored keys and unseen keys.")
    print("  Sweep the readout threshold. Watch the tradeoff between")
    print("  accepting known items and rejecting unseen ones.")
    print()
    K = 30
    for threshold in [0.02, 0.04, 0.06, 0.08, 0.10, 0.15]:
        mem = FourAxisMemory(d=2048, threshold=threshold)
        for i in range(K):
            mem.store("default", f"known_{i}")
        accepted = 0
        for i in range(K):
            r = mem.query("default", f"known_{i}")
            if r["verdict"] == "YES":
                accepted += 1
        rejected = 0
        for i in range(K):
            r = mem.query("default", f"unseen_{i}")
            if r["verdict"] == "UNKNOWN":
                rejected += 1
        print(f"  thr={threshold:>5.2f}  known accepted: "
              f"{accepted:>3}/{K}  unseen rejected: {rejected:>3}/{K}")


# ============================================================
# AXIS 4. Consolidation
# ============================================================
def axis4_consolidation():
    section("AXIS 4. CONSOLIDATION RATE (forgetting)")
    print("  Store 50 items with varying weights. Run 5 consolidation")
    print("  cycles. Anchors (top 20% by weight) are preserved. Others")
    print("  decay by consolidation_rate per cycle.")
    print()
    for rate in [0.0, 0.2, 0.5, 0.8]:
        mem = FourAxisMemory(d=2048, consolidation_rate=rate,
                              anchor_quantile=0.8)
        rng = np.random.default_rng(0)
        for i in range(50):
            w = float(rng.uniform(0.3, 1.0))
            mem.store("default", f"item_{i}", weight=w)
        stats_before = mem.stats()["regions"]["default"]
        for _ in range(5):
            mem.consolidate(cycles=1)
        stats_after = mem.stats()["regions"]["default"]
        # Count retrievable after
        hits_after = 0
        for i in range(50):
            r = mem.query("default", f"item_{i}")
            if r["verdict"] == "YES":
                hits_after += 1
        print(f"  rate={rate:>4.1f}  "
              f"mag_before={stats_before['magnitude_yes']:>6.2f}  "
              f"mag_after={stats_after['magnitude_yes']:>6.2f}  "
              f"retrievable={hits_after}/50")


# ============================================================
# COMBINED. Four profiles
# ============================================================
def combined_profiles():
    section("COMBINED. FOUR MEMORY PROFILES")
    print("  Same code, four different (d, kind, threshold, rate)")
    print("  profiles. Each profile is a distinct memory behavior.")
    print()

    profiles = [
        ("working",   dict(d=512,  threshold=0.10, consolidation_rate=0.8)),
        ("longterm",  dict(d=2048, threshold=0.04, consolidation_rate=0.0)),
        ("seqmem",    dict(d=2048, threshold=0.08, consolidation_rate=0.3)),
        ("contra",    dict(d=2048, threshold=0.05, consolidation_rate=0.1)),
    ]

    for name, params in profiles:
        mem = FourAxisMemory(**params)
        mem.add_region("items", kind="item")
        if name == "contra":
            mem.add_region("claims", kind="anti")
            mem.store("claims", "c1")
            mem.deny("claims", "c1")
            mem.store("claims", "c2")
        if name == "seqmem":
            mem.add_region("events", kind="chain")
            mem.chain("events", ["a", "b", "c", "d", "e"])

        for i in range(20):
            mem.store("items", f"item_{i}")

        stats = mem.stats()
        d_val = stats["d"]
        thr = stats["threshold"]
        rate = stats["consolidation_rate"]
        print(f"  profile '{name}':")
        print(f"    d={d_val}, threshold={thr}, "
              f"consolidation_rate={rate}")

        # Query a few items
        ok = 0
        for i in range(20):
            if mem.query("items", f"item_{i}")["verdict"] == "YES":
                ok += 1
        print(f"    item hits: {ok}/20")

        if name == "seqmem":
            chain = mem.walk("events", "a")
            print(f"    chain walk: {' -> '.join(chain)}")

        if name == "contra":
            r = mem.query("claims", "c1")
            print(f"    contradiction verdict: {r['verdict']}")
        print()


# ============================================================
# Main
# ============================================================
def run_demo(output_dir: str = "holo_four_axis_out") -> None:
    os.makedirs(output_dir, exist_ok=True)
    print("=" * 78)
    print("HOLO-FOUR-AXIS")
    print("A memory store parameterized on four orthogonal axes.")
    print("=" * 78)

    experiments = [
        axis1_dimension,
        axis2_algebra,
        axis3_threshold,
        axis4_consolidation,
        combined_profiles,
    ]
    for f in experiments:
        try:
            f()
        except Exception as e:
            print(f"  EXCEPTION in {f.__name__}: {e}")

    section("OUTPUT")
    demo = FourAxisMemory(d=2048, threshold=0.05,
                          consolidation_rate=0.3)
    demo.add_region("items", kind="item")
    demo.add_region("claims", kind="anti")
    demo.add_region("events", kind="chain")
    for i in range(10):
        demo.store("items", f"fact_{i}", weight=0.5 + 0.05 * i)
    demo.store("claims", "sky_blue")
    demo.deny("claims", "sky_green")
    demo.store("claims", "moon_cheese")
    demo.deny("claims", "moon_cheese")
    demo.chain("events", ["start", "middle", "end"])
    demo.save_json(os.path.join(output_dir, "state.json"))
    demo.plot(os.path.join(output_dir, "regions.png"))

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


def main() -> None:
    parser = argparse.ArgumentParser(
        description="A memory store parameterized on four orthogonal axes.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python holo_four_axis.py
    python holo_four_axis.py --output results/

Library usage:

    from holo_four_axis import FourAxisMemory

    mem = FourAxisMemory(d=2048, threshold=0.05,
                         consolidation_rate=0.3)
    mem.add_region("items", kind="item")
    mem.add_region("events", kind="chain")
    mem.add_region("claims", kind="anti")

    mem.store("items", "apple")
    mem.chain("events", ["wake", "coffee", "work"])
    mem.store("claims", "earth_round")
    mem.deny("claims", "earth_flat")

    for r in ["items", "events", "claims"]:
        print(mem.stats()["regions"][r])
        """)
    parser.add_argument("--output", default="holo_four_axis_out",
                        help="output directory (default: holo_four_axis_out)")
    parser.add_argument("--demo", action="store_true",
                        help="run demo (default behavior)")
    args = parser.parse_args()

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

    run_demo(output_dir=args.output)


if __name__ == "__main__":
    main()