#!/usr/bin/env python3
"""
OLM v3 — Oscillator Language Model with TRAINED oscillators.

v1/v2 kept the oscillator bank frozen random and only fit a linear
readout. v3 trains the physics itself: each hidden unit is a damped
driven oscillator (position h, velocity z) and gradient descent shapes
its frequencies (gamma), damping (eps), couplings (W, Z) and drive (U).
This is the coRNN idea (Rusch & Mishra 2021): oscillator dynamics as
the sequence engine instead of attention.

Still ZERO pre-training: it learns from one novel, from scratch.
Pure NumPy. Hand-derived backprop-through-time. ~150k params at H=256.

"LLM function" comes from the data format: the book is turned into
    ?keyword;original sentence\n
lines, so the model learns that the keyword after '?' steers what the
sentence after ';' is about. At chat time we pick the rarest known word
of the user's message as the keyword — prompt-conditioned generation.

    python3 olm.py --train --save olm.npz      # ~10 min on CPU
    python3 olm.py --load olm.npz              # chat REPL
"""

import argparse
import re
import sys
import time
from collections import Counter
import numpy as np

CHARS = "\n abcdefghijklmnopqrstuvwxyz.,;'-?!\""
C2I = {c: i for i, c in enumerate(CHARS)}
V = len(CHARS)
F32 = np.float32


# ---------------------------------------------------------------- data
def clean(text: str) -> str:
    text = text.replace("’", "'").replace("‘", "'").lower()
    text = text.replace("\r", " ").replace("\n", " ").replace("\t", " ")
    text = re.sub(r"[^a-z .,;'\-?!\"]", " ", text)
    return re.sub(r"  +", " ", text)


def strip_gutenberg(raw: str) -> str:
    start = raw.find("Call me Ishmael")
    end = raw.find("*** END OF THE PROJECT")
    if start != -1:
        raw = raw[start: end if end != -1 else None]
    return re.sub(r"^CHAPTER .*$", " ", raw, flags=re.MULTILINE)


def build_corpus(path: str, seed: int = 0):
    """Turn the novel into '?keyword;sentence\\n' lines + word frequencies."""
    text = clean(strip_gutenberg(open(path, encoding="utf-8", errors="ignore").read()))
    sentences = [s.strip() for s in re.split(r"(?<=[.?!])", text)]
    sentences = [s for s in sentences if 4 <= len(s.split()) <= 45]
    freq = Counter(w for s in sentences for w in re.findall(r"[a-z']+", s))

    lines = []
    kwc = Counter()
    for s in sentences:
        words = [w for w in re.findall(r"[a-z']+", s) if len(w) >= 3 and freq[w] >= 8]
        if not words:
            continue
        kw = min(words, key=lambda w: freq[w])   # rarest usable word = topic
        kwc[kw] += 1
        lines.append(f"?{kw};{s}\n")
    rng = np.random.default_rng(seed)
    rng.shuffle(lines)
    return "".join(lines), freq, kwc


# ---------------------------------------------------------------- model
def init_params(H: int, seed: int = 0):
    rng = np.random.default_rng(seed)
    s = 1.0 / np.sqrt(H)
    return {
        "W":  rng.normal(0, s, (H, H)).astype(F32),   # position coupling
        "Z":  rng.normal(0, s, (H, H)).astype(F32),   # velocity coupling
        "U":  rng.normal(0, 0.5, (V, H)).astype(F32), # char drive
        "b":  np.zeros(H, F32),
        "gamma": np.ones(H, F32),                     # restoring force (freq^2)
        "eps":   np.ones(H, F32),                     # damping
        "Wo": rng.normal(0, s, (H, V)).astype(F32),
        "bo": np.zeros(V, F32),
    }


def forward(p, h, z, X, Y, dt):
    """One BPTT window. X,Y: (B,T) int ids. Returns loss, caches, final state."""
    B, T = X.shape
    caches, loss = [], 0.0
    for t in range(T):
        h0, z0 = h, z
        a = h0 @ p["W"] + z0 @ p["Z"] + p["U"][X[:, t]] + p["b"]
        c = np.tanh(a)
        z = z0 + dt * (c - p["gamma"] * h0 - p["eps"] * z0)
        h = h0 + dt * z
        logits = h @ 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((h0, z0, c, probs, X[:, t], Y[:, t], h))
    return loss / (B * T), caches, h, z


def backward(p, caches, dt):
    B = caches[0][0].shape[0]
    T = len(caches)
    g = {k: np.zeros_like(v) for k, v in p.items()}
    gh = np.zeros_like(caches[0][0])
    gz = np.zeros_like(gh)
    for t in reversed(range(T)):
        h0, z0, c, probs, x, y, h = caches[t]
        # output head
        dlog = probs.copy()
        dlog[np.arange(B), y] -= 1.0
        dlog /= (B * T)
        g["Wo"] += h.T @ dlog
        g["bo"] += dlog.sum(0)
        gh = gh + dlog @ p["Wo"].T
        # h = h0 + dt*z
        gz = gz + dt * gh
        # z = z0 + dt*(c - gamma*h0 - eps*z0)
        gc = dt * gz
        g["gamma"] += (-dt * gz * h0).sum(0)
        g["eps"]   += (-dt * gz * z0).sum(0)
        gh_prev = gh - dt * (p["gamma"] * gz)
        gz_prev = gz * (1.0 - dt * p["eps"])
        # a = h0@W + z0@Z + U[x] + b ; c = tanh(a)
        ga = gc * (1.0 - c * c)
        g["W"] += h0.T @ ga
        g["Z"] += z0.T @ ga
        np.add.at(g["U"], x, ga)
        g["b"] += ga.sum(0)
        gh = gh_prev + ga @ p["W"].T
        gz = gz_prev + ga @ p["Z"].T
    return g


def gradcheck():
    """Finite-difference check of the hand-derived BPTT gradients."""
    np.random.seed(1)
    H, B, T, dt = 6, 2, 4, 0.5
    p = {k: v.astype(np.float64) for k, v in init_params(H, seed=3).items()}
    X = np.random.randint(0, V, (B, T))
    Y = np.random.randint(0, V, (B, T))
    h = np.random.randn(B, H) * 0.1
    z = np.random.randn(B, H) * 0.1
    loss, caches, _, _ = forward(p, h, z, 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(5, flat.size), replace=False):
            eps_ = 1e-5
            old = flat[idx]
            flat[idx] = old + eps_
            lp, _, _, _ = forward(p, h, z, X, Y, dt)
            flat[idx] = old - eps_
            lm, _, _, _ = forward(p, h, z, 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)
    # keep the physics stable: restoring force and damping stay positive
    np.clip(p["gamma"], 0.05, 4.0, out=p["gamma"])
    np.clip(p["eps"], 0.05, 4.0, out=p["eps"])
    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, dt = args.batch, args.seq, args.hidden, args.dt

    p = init_params(H)
    n_params = sum(x.size for x in p.values())
    print(f"OLM v3 | hidden {H} | 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)
        h = np.zeros((B, H), F32)
        z = np.zeros((B, H), F32)
        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, h, z = forward(p, h, z, 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, H)
        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, H, B=50):
    L = len(ids) // B
    st = ids[: B * L].reshape(B, L)
    h = np.zeros((B, H), F32)
    z = np.zeros((B, H), F32)
    loss, caches, _, _ = forward(p, h, z, st[:, :-1][:, :400], st[:, 1:][:, :400], dt)
    probs_hits = 0
    tot = 0
    for (h0, z0, c, probs, x, y, hh) in caches[50:]:   # skip washout
        probs_hits += int((probs.argmax(1) == y).sum())
        tot += len(y)
    return loss, probs_hits / tot


# ---------------------------------------------------------------- generation
def step_one(p, h, z, cid, dt):
    a = h @ p["W"] + z @ p["Z"] + p["U"][cid] + p["b"]
    c = np.tanh(a)
    z = z + dt * (c - p["gamma"] * h - p["eps"] * z)
    h = h + dt * z
    return h, z, h @ 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 = p["b"].size
    h = np.zeros((1, H), F32)
    z = np.zeros((1, H), F32)
    logits = None
    for ch in prime:
        h, z, logits = step_one(p, h, z, np.array([C2I.get(ch, 1)]), dt)
    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])
        h, z, logits = step_one(p, h, z, np.array([nxt]), dt)
    return "".join(out)


STOPWORDS = set("""the and that this with what who how which but not nor all any
are was were been being have has had having you your yours thou thee thy they
them their there then than when where why while will would could should shall
may might must can does did doing about into from upon over under very much
more most some such only just also even ever never here his her him its our
ours out off for one two too tell know said says say see seem like want need
good well make made take give get got come came went gone let yes now""".split())


def pick_keyword(msg, freq, kwc):
    words = [w for w in re.findall(r"[a-z']+", msg.lower())
             if len(w) >= 3 and w not in STOPWORDS and freq.get(w, 0) >= 8]
    if not words:
        return "whale"
    trained = [w for w in words if kwc.get(w, 0) >= 3]
    # most specific topic the model actually practised on; else rarest known
    return min(trained or words, key=lambda w: freq[w])


# ---------------------------------------------------------------- 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="OLM v3 — trained oscillators")
    ap.add_argument("--text", default="mobydick.txt")
    ap.add_argument("--hidden", type=int, default=256)
    ap.add_argument("--seq", type=int, default=128)
    ap.add_argument("--batch", type=int, default=64)
    ap.add_argument("--dt", type=float, default=0.5)
    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="olm.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}\nolm ({kw}) : {resp}")
        return
    print("OLM 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"olm ({kw}) > {resp}")


if __name__ == "__main__":
    main()