#!/usr/bin/env python3
"""gary-neuron-chat: a tiny dialogue model that LEARNS FROM USE on three timescales.
  FAST (within a conversation): a plastic hippocampal memory (softmax attention over
       the conversation) recalls facts you stated earlier. No weights move.
  SLOW (across sessions): `sleep` replays buffered conversations mixed with base
       corpus (experience replay) and fine-tunes the cortex -- learning without forgetting.
  The brain (cortex weights + replay buffer + age) persists in brain.npz and evolves.
Cortex = gary-4-petite fine-tuned on SODA dialogue. Plastic memory meta-trained
to retrieve. Pure numpy, no torch.  Usage: python brain.py [chat|demo|sleep]"""
import os, sys, json, numpy as np
import gpt_numpy as G
from tokenizers import ByteLevelBPETokenizer
D=os.path.dirname(os.path.abspath(__file__))
CFG=dict(E=96,H=4,L=4,BLK=128); EOT=0
tok=ByteLevelBPETokenizer(f"{D}/petite_vocab.json", f"{D}/petite_merges.txt")
def _f(x): return float(np.ravel(x)[0])

def load_brain(path=f"{D}/brain.npz"):
    if os.path.exists(path):
        z=np.load(path,allow_pickle=True)
        P={k[2:]:z[k].astype(np.float32) for k in z.files if k.startswith("P/")}
        Hm={k[3:]:z[k] for k in z.files if k.startswith("hm/")}
        buf=list(z["buffer"]) if "buffer" in z.files else []; age=int(z["age"]) if "age" in z.files else 0
    else:
        zc=np.load(f"{D}/cortex.npz",allow_pickle=True); P={k[2:]:zc[k].astype(np.float32) for k in zc.files if k.startswith("P/")}
        zh=np.load(f"{D}/hippo.npz"); Hm={k:zh[k] for k in zh.files}; buf=[]; age=0
    return P,Hm,buf,age

def save_brain(P,Hm,buf,age,path=f"{D}/brain.npz"):
    d={"P/"+k:v for k,v in P.items()}; d.update({"hm/"+k:np.atleast_1d(v) for k,v in Hm.items()})
    d["buffer"]=np.array(buf[-400:],dtype=object); d["age"]=age; np.savez(path,**d)

def cortex(P, ids):
    x=np.array([ids[-CFG["BLK"]:]]); logits,cache=G.forward(P,x,CFG); return cache["xf"][0], logits[0]

def hippo_bias(Hm, xf, logits, win_ids):
    Wq=Hm["Wq"].astype(np.float64); bN=_f(Hm["bN"]); rdec=_f(Hm["rdec"]); g=_f(Hm["g"]); T=xf.shape[0]
    if T<2: return np.zeros(logits.shape[1])
    Q=xf.astype(np.float64)@Wq; qa=Q[-1]; sim=Q@qa
    lp=logits-logits.max(1,keepdims=True); lp=lp-np.log(np.exp(lp).sum(1,keepdims=True))
    nxt=np.array(list(win_ids[1:])+[0]); nll=np.zeros(T); nll[:-1]=-lp[np.arange(T-1),nxt[:-1]]
    t=np.arange(T); dist=(T-1)-t; valid=t<(T-1)
    score=np.where(valid, sim+bN*nll+rdec*dist, -1e9); score-=score.max(); a=np.exp(score)*valid; a/=a.sum()+1e-12
    S=np.zeros(logits.shape[1]); np.add.at(S, nxt, a); return g*S

def gen_reply(P,Hm,ids,max_new=20,temp=0.0,use_hippo=True,seed=0):
    rng=np.random.default_rng(seed); out=[]
    for _ in range(max_new):
        win=ids[-CFG["BLK"]:]; xf,logits=cortex(P,win); ll=logits[-1].astype(np.float64)
        if use_hippo: ll=ll+hippo_bias(Hm,xf,logits,win)
        if temp<=0: nx=int(ll.argmax())
        else:
            q=np.exp(ll/temp-(ll/temp).max()); q/=q.sum(); nx=int(rng.choice(len(q),p=q))
        if nx==EOT: break
        ids=ids+[nx]; out.append(nx)
        if tok.decode(out).endswith("\n"): break
    return tok.decode(out).strip(), ids

def sleep(P, buf, base_path=f"{D}/train.bin", steps=30, lr=1.2e-4, mix=0.12):
    """Consolidate: replay buffered convos + base corpus (mostly base -> no forgetting)."""
    if not buf: return P, None
    base=np.memmap(base_path,dtype=np.uint16,mode="r"); T=CFG["BLK"]
    ci=[]; [ci.extend(tok.encode(t).ids+[EOT]) for t in buf]; convo=np.array(ci,dtype=np.int64)
    opt=G.Adam(P,lr=lr); rng=np.random.default_rng(0)
    for _ in range(steps):
        xs=[];ys=[]
        for _ in range(16):
            if rng.random()<mix and len(convo)>T+1: s=convo[(i:=rng.integers(0,len(convo)-T-1)):i+T+1]
            else: j=rng.integers(0,len(base)-T-1); s=np.asarray(base[j:j+T+1],dtype=np.int64)
            xs.append(s[:T]); ys.append(s[1:T+1])
        loss,cache=G.forward(P,np.stack(xs),CFG,np.stack(ys)); opt.step(P,G.backward(P,CFG,cache))
    return P, None

def chat():
    P,Hm,buf,age=load_brain(); print(f"gary-neuron-chat (age {age}). Tell me about yourself; I'll remember within our chat. '/sleep' to consolidate, '/exit' to leave.")
    ids=tok.encode("G: hi").ids; sess=[]
    while True:
        try: u=input("\nyou: ").strip()
        except (EOFError,KeyboardInterrupt): break
        if u=="/exit": break
        if u=="/sleep":
            P,_=sleep(P,sess); age+=1; save_brain(P,Hm,buf+sess,age); sess=[]; print("...slept; cortex consolidated."); continue
        ids=ids+tok.encode(f"\nU: {u}\nG:").ids; r,ids=gen_reply(P,Hm,ids); print("gary:",r); sess.append(f"U: {u}\nG: {r}")
    save_brain(P,Hm,buf+sess,age); print("brain saved.")

if __name__=="__main__":
    cmd=sys.argv[1] if len(sys.argv)>1 else "chat"
    if cmd=="chat": chat()
    elif cmd=="demo":
        P,Hm,buf,age=load_brain()
        ids=tok.encode("G: hello").ids
        for u in ["my dog is named Buddy","how is the weather","i like to read books","what time is it"]:
            ids=ids+tok.encode(f"\nU: {u}\nG: ok").ids
        q="what is my dog 's name ?"; qids=ids+tok.encode(f"\nU: {q}\nG:").ids
        on,_=gen_reply(P,Hm,qids,use_hippo=True); off,_=gen_reply(P,Hm,qids,use_hippo=False)
        print(f"Q: {q}\n  plastic-memory ON : {on!r}\n  plastic-memory OFF: {off!r}")