# Experiments for "Modern Hopfield Networks": dense associative memory (polynomial and exponential interaction),
# the continuous modern Hopfield network (softmax update), its capacity, energy decrease and convergence. Writes modern.json.
import json, math, numpy as np
rng = np.random.default_rng(0); out = {}

# ---------- 1) binary dense associative memory (Krotov & Hopfield 2016; Demircigil et al. 2017) ----------
# E(s) = -sum_mu F(x_mu . s). Update neuron i: compare the energy with s_i = +1 and s_i = -1, keep the lower one.
def F(z, kind):
    if kind == "exp": return np.exp(z - z.max())          # shift for stability: the comparison only needs ratios
    return z.astype(float) ** kind
def recall_dense(X, s, kind, rng, sweeps=20):
    s = s.copy(); d = len(s); m = X @ s                    # m_mu = x_mu . s, updated incrementally
    for _ in range(sweeps):
        ch = 0
        for i in rng.permutation(d):
            m_minus = m - X[:, i] * s[i]                     # overlaps without neuron i
            if kind == "exp":
                a, b = m_minus + X[:, i], m_minus - X[:, i]; top = max(a.max(), b.max())
                plus = np.exp(a - top).sum(); minus = np.exp(b - top).sum()
            else:
                plus = F(m_minus + X[:, i], kind).sum(); minus = F(m_minus - X[:, i], kind).sum()
            v = 1.0 if plus >= minus else -1.0
            if v != s[i]: m = m_minus + X[:, i] * v; s[i] = v; ch += 1
        if not ch: break
    return s
def capacity_curve(kind, d, Ns, noise=0.1, trials=40):
    res = []
    for N in Ns:
        ok = 0
        for t in range(trials):
            X = rng.choice([-1.0, 1.0], size=(N, d)); mu = rng.integers(N)
            q = X[mu].copy(); q[rng.choice(d, int(noise * d), replace=False)] *= -1
            ok += np.array_equal(recall_dense(X, q, kind, rng), X[mu])
        res.append({"N": int(N), "ok": ok / trials})
        if ok / trials < 0.3: break
    return res
d = 64; out["dense"] = {}
for kind, Ns in [(2, [4, 6, 8, 10, 12, 16]), (3, [16, 32, 64, 96, 128, 192, 256]), (4, [64, 128, 256, 512, 1024, 2048]), ("exp", [256, 1024, 4096, 16384])]:
    out["dense"][str(kind)] = capacity_curve(kind, d, Ns); print("dense", kind, out["dense"][str(kind)], flush=True)

# ---------- 2) the 8 shapes again ----------
S = json.load(open("shapes.json")); P = np.array([[1.0 if c == "1" else -1.0 for c in p] for p in S["patterns"]])
shapes = {}
for kind in (2, 3, 4, "exp"):
    ok = 0; tries = 0
    for mu in range(8):
        for t in range(20):
            q = P[mu].copy(); q[rng.choice(256, 38, replace=False)] *= -1
            ok += np.array_equal(recall_dense(P, q, kind, rng), P[mu]); tries += 1
    shapes[str(kind)] = ok / tries; print("8 shapes, F =", kind, ok / tries)
# continuous modern Hopfield on the shapes: one update xi <- X softmax(beta X^T xi), read out by sign
def softmax(z): z = z - z.max(); e = np.exp(z); return e / e.sum()
for beta in (0.01, 0.05, 0.1, 1.0):
    ok = 0; tries = 0
    for mu in range(8):
        for t in range(20):
            q = P[mu].copy(); q[rng.choice(256, 38, replace=False)] *= -1
            xi = P.T @ softmax(beta * P @ q); ok += np.array_equal(np.sign(xi), P[mu]); tries += 1
    shapes[f"softmax_beta_{beta}"] = ok / tries; print("8 shapes, one softmax step, beta =", beta, ok / tries)
out["shapes"] = shapes

# ---------- 3) continuous modern Hopfield: capacity vs dimension (one step, 10% of bits flipped, beta = 1) ----------
cap = []
for d in (10, 20, 30, 40):                     # 10% of the bits flipped exactly: 1, 2, 3, 4
    best = 0; rows = []
    for N in sorted({int(round(2 ** (k / 2))) for k in range(2, 40)}):
        ok = 0; trials = 60
        for t in range(trials):
            X = rng.choice([-1.0, 1.0], size=(N, d)); mu = rng.integers(N)
            q = X[mu].copy(); q[rng.choice(d, d // 10, replace=False)] *= -1
            xi = X.T @ softmax(1.0 * X @ q); ok += np.array_equal(np.sign(xi), X[mu])
        rows.append({"N": N, "ok": ok / trials})
        if ok / trials >= 0.9: best = N
        if ok / trials < 0.5 or N * d > 4e7: break
    cap.append({"d": d, "N90": best, "rows": rows}); print("softmax capacity d =", d, "largest N with >= 90% one-step recall:", best, flush=True)
# classical, same protocol (asynchronous recall until stable)
for d in (10, 20, 30, 40):
    best = 0
    for N in range(1, 16):
        ok = 0
        for t in range(60):
            X = rng.choice([-1.0, 1.0], size=(N, d)); mu = rng.integers(N)
            q = X[mu].copy(); q[rng.choice(d, d // 10, replace=False)] *= -1
            ok += np.array_equal(recall_dense(X, q, 2, rng), X[mu])
        if ok / 60 >= 0.9: best = N
    next(c for c in cap if c["d"] == d)["classical_N90"] = best
print("classical:", [(c["d"], c["classical_N90"]) for c in cap])
out["capacity"] = cap

# ---------- 4) the continuous energy never increases under the update (CCCP), and one step is usually enough ----------
def energy(X, xi, beta):     # E = -lse(beta, X^T xi) + 1/2 xi^T xi (constants dropped); X: (N, d)
    z = beta * X @ xi; m = z.max(); return -(m + math.log(np.exp(z - m).sum())) / beta + 0.5 * xi @ xi
viol = 0; steps = []
for t in range(300):
    d, N = 32, 50; X = rng.normal(size=(N, d)); beta = float(rng.choice([0.05, 0.2, 1.0]))
    xi = rng.normal(size=d) * 2
    for k in range(100):
        new = X.T @ softmax(beta * X @ xi)
        viol += energy(X, new, beta) > energy(X, xi, beta) + 1e-9
        if np.linalg.norm(new - xi) < 1e-8: break
        xi = new
    steps.append((beta, k))
out["energy_increases"] = int(viol)
out["steps_to_converge"] = {str(b): float(np.median([k for bb, k in steps if bb == b])) for b in (0.05, 0.2, 1.0)}
print("energy increases:", viol, "median steps to converge per beta:", out["steps_to_converge"])
# one-step retrieval for well separated patterns: start near pattern 0, error after 1 step vs beta
sep = []
d, N = 64, 100; X = rng.choice([-1.0, 1.0], size=(N, d)) ; q = X[0].copy(); q[rng.choice(d, 6, replace=False)] *= -1
for beta in (0.02, 0.05, 0.1, 0.2, 0.5, 1.0):
    xi = X.T @ softmax(beta * X @ q); sep.append({"beta": beta, "err": float(np.linalg.norm(xi - X[0]) / np.linalg.norm(X[0])), "p0": float(softmax(beta * X @ q)[0])})
out["one_step_error"] = sep; print("one-step error vs beta:", [(r["beta"], round(r["err"], 4), round(r["p0"], 3)) for r in sep])
# ---------- 5) the three regimes: 4 clusters of 5 similar patterns, a noisy copy of pattern 7 (cluster 1) ----------
rg = np.random.default_rng(1); d = 64; cent = rg.normal(size=(4, d))
X = np.stack([cent[c] + 0.33 * rg.normal(size=d) for c in range(4) for _ in range(5)]); X = X / np.linalg.norm(X, axis=1, keepdims=True) * math.sqrt(d)
q0 = X[7] + 0.5 * rg.normal(size=d); regimes = []
for beta in (0.001, 0.05, 0.2, 1.0):
    xi = q0.copy()
    for k in range(500):
        new = X.T @ softmax(beta * X @ xi)
        if np.linalg.norm(new - xi) < 1e-7: break
        xi = new
    p = softmax(beta * X @ xi)
    regimes.append({"beta": beta, "cluster_weight": [round(float(p[5 * c:5 * c + 5].sum()), 3) for c in range(4)], "max_weight": round(float(p.max()), 3), "argmax": int(p.argmax()), "steps": k})
    print("regime", regimes[-1])
out["regimes"] = regimes
json.dump(out, open("modern.json", "w"), indent=1)