#!/usr/bin/env python3
"""
KOLM — Kuramoto Oscillator Language Model (OLM v4).

v3 (olm.py) trains damped *amplitude* oscillators (coRNN). v4 swaps the core
for Kuramoto oscillators, AKOrN-style (Miyato et al., ICLR 2025): each hidden
unit is a unit vector x_i on the sphere S^{N-1}. Units interact through one
trained scalar coupling J_ij per pair and synchronize; binding-by-synchrony
means units that align are bound into one active concept, and the same units
re-align differently for another concept.

One char step (all trained: J, A, U, Wo, bo):

    y_i = sum_j J_ij x_j + U[char]_i              coupling + input drive
    P_i = y_i - (x_i . y_i) x_i                   project onto tangent space
    x_i <- normalize(x_i + dt (A_i x_i + P_i))    A_i antisymmetric: rotation

Unit norm is enforced by construction, so stability is geometric — no
gamma/eps clipping. Same zero-pretraining setup as v3: learns one novel as
'?keyword;sentence' lines; chat picks a keyword and generates conditioned
on it. Pure NumPy, hand-derived BPTT, ~141k params at H=256 N=4.

    python3 kolm.py --gradcheck
    python3 kolm.py --train --save kolm256.npz
    python3 kolm.py --load kolm256.npz
"""

import argparse
import sys
import time
import numpy as np

from olm import CHARS, C2I, V, F32, build_corpus, pick_keyword


# ---------------------------------------------------------------- model
def init_params(H: int, N: int, seed: int = 0, om_init: float = 0.1,
                groups: int = 0):
    rng = np.random.default_rng(seed)
    p = {
        "J":  rng.normal(0, 1.0 / np.sqrt(H), (H, H)).astype(F32),
        "Om": rng.normal(0, om_init, (H, N, N)).astype(F32),  # A = Om - Om^T
        "U":  rng.normal(0, 0.5, (V, H, N)).astype(F32),
        "Wo": rng.normal(0, 1.0 / np.sqrt(H * N),
                         (H * N + groups, V)).astype(F32),
        "bo": np.zeros(V, F32),
    }
    if groups:
        # coherence readout: r_g = ||sum_i M_gi x_i|| — a trained order
        # parameter, large only when the units row g weights are in sync
        p["M"] = rng.normal(0, 1.0 / np.sqrt(H), (groups, H)).astype(F32)
    return p


def init_state(B: int, H: int, N: int, dtype=F32):
    x = np.zeros((B, H, N), dtype)
    x[..., 0] = 1.0        # all units start at the same pole; A and U disperse them
    return x


def antisym(Om):
    return Om - Om.transpose(0, 2, 1)


def couple(J, x):
    """y[b,i,:] = sum_j J[i,j] x[b,j,:] as one BLAS matmul."""
    B, H, N = x.shape
    return (J @ x.transpose(1, 0, 2).reshape(H, B * N)) \
        .reshape(-1, B, N).transpose(1, 0, 2)


def readout_feats(p, x):
    """[flat phases, coherence r] and the population vectors m (or None)."""
    B, H, N = x.shape
    if "M" not in p:
        return x.reshape(B, H * N), None, None
    m = couple(p["M"], x)                                # (B,G,N)
    r = np.sqrt((m * m).sum(-1) + 1e-8)                  # (B,G)
    return np.concatenate([x.reshape(B, H * N), r], axis=1), m, r


def forward(p, x, X, Y, dt):
    """One BPTT window. X,Y: (B,T) int ids. Returns loss, caches, final state."""
    B, T = X.shape
    H, N = p["Om"].shape[0], p["Om"].shape[1]
    A = antisym(p["Om"])
    caches, loss = [], 0.0
    for t in range(T):
        x0 = x
        y = couple(p["J"], x0) + p["U"][X[:, t]]
        d = (x0 * y).sum(-1, keepdims=True)
        P = y - d * x0
        R = np.einsum("inm,bim->bin", A, x0)
        xt = x0 + dt * (R + P)
        nrm = np.sqrt((xt * xt).sum(-1, keepdims=True))
        x = xt / nrm
        f, m, r = readout_feats(p, x)
        logits = f @ p["Wo"] + p["bo"]
        logits -= logits.max(axis=1, keepdims=True)
        e = np.exp(logits)
        probs = e / e.sum(axis=1, keepdims=True)
        loss -= np.log(probs[np.arange(B), Y[:, t]] + 1e-9).sum()
        caches.append((x0, y, d, nrm, x, f, m, r, probs, X[:, t], Y[:, t]))
    return loss / (B * T), caches, x


def backward(p, caches, dt):
    B = caches[0][0].shape[0]
    T = len(caches)
    H, N = p["Om"].shape[0], p["Om"].shape[1]
    A = antisym(p["Om"])
    g = {k: np.zeros_like(v) for k, v in p.items()}
    gA = np.zeros_like(A)
    gx = np.zeros_like(caches[0][0])
    for t in reversed(range(T)):
        x0, y, d, nrm, x1, f, m, r, probs, xs, ys = caches[t]
        # output head
        dlog = probs.copy()
        dlog[np.arange(B), ys] -= 1.0
        dlog /= (B * T)
        g["Wo"] += f.T @ dlog
        g["bo"] += dlog.sum(0)
        gf = dlog @ p["Wo"].T
        gx1 = gx + gf[:, :H * N].reshape(B, H, N)
        if m is not None:
            # r = ||m||, m_g = sum_i M_gi x1_i
            gm = (gf[:, H * N:] / r)[:, :, None] * m
            gm_r = gm.transpose(1, 0, 2).reshape(-1, B * N)
            x1_r = x1.transpose(1, 0, 2).reshape(H, B * N)
            g["M"] += gm_r @ x1_r.T
            gx1 += (p["M"].T @ gm_r).reshape(H, B, N).transpose(1, 0, 2)
        # x1 = xt / ||xt||  =>  gxt = (I - x1 x1^T) gx1 / ||xt||
        gxt = (gx1 - (gx1 * x1).sum(-1, keepdims=True) * x1) / nrm
        # xt = x0 + dt*(R + P)
        gx0 = gxt.copy()
        gR = dt * gxt
        gP = dt * gxt
        # R_i = A_i x0_i
        gA += np.einsum("bin,bim->inm", gR, x0)
        gx0 += np.einsum("inm,bin->bim", A, gR)
        # P = y - (x0.y) x0
        gPdot = (gP * x0).sum(-1, keepdims=True)
        gy = gP - gPdot * x0
        gx0 -= gPdot * y + d * gP
        # y = J x0 + U[xs]
        gy_r = gy.transpose(1, 0, 2).reshape(H, B * N)
        x0_r = x0.transpose(1, 0, 2).reshape(H, B * N)
        g["J"] += gy_r @ x0_r.T
        gx0 += (p["J"].T @ gy_r).reshape(H, B, N).transpose(1, 0, 2)
        np.add.at(g["U"], xs, gy)
        gx = gx0
    g["Om"] = gA - gA.transpose(0, 2, 1)
    return g


def gradcheck():
    """Finite-difference check of the hand-derived BPTT gradients."""
    np.random.seed(1)
    H, N, B, T, dt = 5, 3, 2, 4, 0.5
    p = {k: v.astype(np.float64)
         for k, v in init_params(H, N, seed=3, groups=3).items()}
    X = np.random.randint(0, V, (B, T))
    Y = np.random.randint(0, V, (B, T))
    x = np.random.randn(B, H, N)
    x /= np.sqrt((x * x).sum(-1, keepdims=True))
    loss, caches, _ = forward(p, x, X, Y, dt)
    g = backward(p, caches, dt)
    worst = 0.0
    for name in p:
        flat = p[name].reshape(-1)
        for idx in np.random.choice(flat.size, min(6, flat.size), replace=False):
            eps_ = 1e-5
            old = flat[idx]
            flat[idx] = old + eps_
            lp, _, _ = forward(p, x, X, Y, dt)
            flat[idx] = old - eps_
            lm, _, _ = forward(p, x, X, Y, dt)
            flat[idx] = old
            num = (lp - lm) / (2 * eps_)
            ana = g[name].reshape(-1)[idx]
            rel = abs(num - ana) / max(1e-8, abs(num) + abs(ana))
            worst = max(worst, rel)
    print(f"gradcheck worst relative error: {worst:.2e} "
          f"({'PASS' if worst < 1e-4 else 'FAIL'})")
    return worst < 1e-4


# ---------------------------------------------------------------- training
def adam_step(p, g, m, v, t, lr, clip=1.0):
    norm = np.sqrt(sum(float((gi ** 2).sum()) for gi in g.values()))
    scale = min(1.0, clip / (norm + 1e-8))
    b1, b2, e = 0.9, 0.999, 1e-8
    for k in p:
        gk = g[k] * scale
        m[k] = b1 * m[k] + (1 - b1) * gk
        v[k] = b2 * v[k] + (1 - b2) * gk * gk
        mh = m[k] / (1 - b1 ** t)
        vh = v[k] / (1 - b2 ** t)
        p[k] -= (lr * mh / (np.sqrt(vh) + e)).astype(F32)
    return norm


def train(args):
    corpus, freq, kwc = build_corpus(args.text)
    ids = np.array([C2I[c] for c in corpus], dtype=np.int64)
    n_val = 20_000
    tr, va = ids[:-n_val], ids[-n_val:]
    B, T, H, N, dt = args.batch, args.seq, args.hidden, args.ndim, args.dt

    p = init_params(H, N, om_init=args.om_init, groups=args.groups)
    n_params = sum(x.size for x in p.values())
    print(f"KOLM v4 | hidden {H} x S^{N - 1} | params {n_params:,} "
          f"({n_params * 4 / 1e6:.2f} MB) | corpus {len(ids):,} chars", flush=True)

    L = len(tr) // B
    streams = tr[: B * L].reshape(B, L)
    m = {k: np.zeros_like(x) for k, x in p.items()}
    v = {k: np.zeros_like(x) for k, x in p.items()}
    step = 0
    for ep in range(1, args.epochs + 1):
        lr = args.lr * (0.5 ** max(0, ep - args.epochs + 6) if ep > args.epochs - 6 else 1.0)
        x = init_state(B, H, N)
        tot = nb = 0
        t0 = time.time()
        for s in range(0, L - T - 1, T):
            X, Y = streams[:, s:s + T], streams[:, s + 1:s + T + 1]
            loss, caches, x = forward(p, x, X, Y, dt)
            g = backward(p, caches, dt)
            step += 1
            adam_step(p, g, m, v, step, lr)
            tot += loss
            nb += 1
        vl, vacc = evaluate(p, va, dt)
        print(f"epoch {ep:2d} | train loss {tot / nb:.3f} | "
              f"val loss {vl:.3f} | val acc {vacc:.1%} | "
              f"{time.time() - t0:.0f}s", flush=True)
        if ep % 4 == 0 or ep == args.epochs:
            print("  sample:", generate(p, dt, "?whale;", seed=ep)[:110], flush=True)
    save(args.save, p, dt, freq, kwc)
    print(f"saved to {args.save}", flush=True)


def evaluate(p, ids, dt, B=50):
    H, N = p["Om"].shape[0], p["Om"].shape[1]
    L = len(ids) // B
    st = ids[: B * L].reshape(B, L)
    x = init_state(B, H, N)
    loss, caches, _ = forward(p, x, st[:, :-1][:, :400], st[:, 1:][:, :400], dt)
    hits = tot = 0
    for cache in caches[50:]:   # skip washout
        probs, ys = cache[-3], cache[-1]
        hits += int((probs.argmax(1) == ys).sum())
        tot += len(ys)
    return loss, hits / tot


# ---------------------------------------------------------------- generation
def step_one(p, x, cid, dt, A):
    x0 = x
    y = couple(p["J"], x0) + p["U"][cid]
    d = (x0 * y).sum(-1, keepdims=True)
    R = np.einsum("inm,bim->bin", A, x0)
    xt = x0 + dt * (R + y - d * x0)
    x = xt / np.sqrt((xt * xt).sum(-1, keepdims=True))
    f, _, _ = readout_feats(p, x)
    return x, f @ p["Wo"] + p["bo"]


def generate(p, dt, prime, max_len=300, temperature=0.8, top_k=8, seed=None):
    rng = np.random.default_rng(seed)
    H, N = p["Om"].shape[0], p["Om"].shape[1]
    A = antisym(p["Om"])
    x = init_state(1, H, N)
    logits = None
    for ch in prime:
        x, logits = step_one(p, x, np.array([C2I.get(ch, 1)]), dt, A)
    out = []
    for _ in range(max_len):
        lg = logits[0] / temperature
        lg -= lg.max()
        pr = np.exp(lg)
        if top_k and top_k < V:
            cut = np.partition(pr, -top_k)[-top_k]
            pr = np.where(pr >= cut, pr, 0.0)
        pr /= pr.sum()
        nxt = int(rng.choice(V, p=pr))
        if CHARS[nxt] == "\n":
            break
        out.append(CHARS[nxt])
        x, logits = step_one(p, x, np.array([nxt]), dt, A)
    return "".join(out)


# ---------------------------------------------------------------- persistence
def save(path, p, dt, freq, kwc):
    words = np.array(list(freq.keys()))
    counts = np.array(list(freq.values()))
    kwords = np.array(list(kwc.keys()))
    kcounts = np.array(list(kwc.values()))
    np.savez_compressed(path, dt=dt, words=words, counts=counts,
                        kwords=kwords, kcounts=kcounts,
                        **{f"p_{k}": x for k, x in p.items()})


def load(path):
    zf = np.load(path)
    p = {k[2:]: zf[k] for k in zf.files if k.startswith("p_")}
    freq = dict(zip(zf["words"].tolist(), zf["counts"].tolist()))
    kwc = dict(zip(zf["kwords"].tolist(), zf["kcounts"].tolist()))
    return p, float(zf["dt"]), freq, kwc


# ---------------------------------------------------------------- main
def main():
    ap = argparse.ArgumentParser(description="KOLM v4 — Kuramoto oscillators")
    ap.add_argument("--text", default="mobydick.txt")
    ap.add_argument("--hidden", type=int, default=256)
    ap.add_argument("--ndim", type=int, default=4, help="oscillator dimension N")
    ap.add_argument("--seq", type=int, default=128)
    ap.add_argument("--batch", type=int, default=64)
    ap.add_argument("--dt", type=float, default=0.25)
    ap.add_argument("--om-init", type=float, default=0.1,
                    help="init scale of Om; with dt sets the gradient horizon")
    ap.add_argument("--groups", type=int, default=32,
                    help="coherence readout groups G (0 disables)")
    ap.add_argument("--lr", type=float, default=2e-3)
    ap.add_argument("--epochs", type=int, default=18)
    ap.add_argument("--train", action="store_true")
    ap.add_argument("--save", default="kolm256.npz")
    ap.add_argument("--load", default=None)
    ap.add_argument("--ask", default=None, help="one-shot question")
    ap.add_argument("--gradcheck", action="store_true")
    args = ap.parse_args()

    if args.gradcheck:
        sys.exit(0 if gradcheck() else 1)
    if args.train:
        train(args)
        if not args.ask:
            return
    p, dt, freq, kwc = load(args.load or args.save)

    def answer(msg):
        kw = pick_keyword(msg, freq, kwc)
        return kw, generate(p, dt, f"?{kw};")

    if args.ask:
        kw, resp = answer(args.ask)
        print(f"you : {args.ask}\nkolm ({kw}) : {resp}")
        return
    print("KOLM chat — it answers about the topic word of your message (ctrl-d quits)")
    while True:
        try:
            msg = input("\nyou > ")
        except EOFError:
            break
        kw, resp = answer(msg)
        print(f"kolm ({kw}) > {resp}")


if __name__ == "__main__":
    main()