"""gary-neuron-chat benchmark: (1) fact-recall with the plastic memory ON vs OFF,
(2) sleep consolidation learns new material while base English stays intact."""
import os, numpy as np
import gpt_numpy as G
from brain import load_brain, sleep, CFG, tok, EOT
import hippo
D=os.path.dirname(os.path.abspath(__file__))
P,Hm,buf,age=load_brain()

# ---- 1. fact recall (held-out episodes, plastic memory ON vs OFF) ----
va=dict(np.load(f"{D}/episodes_val.npz")); on,off=hippo.evaluate(Hm,va)
print(f"[fact recall, {len(va['ids'])} held-out episodes]")
print(f"   plastic memory ON : {on*100:.1f}%   OFF (frozen cortex alone): {off*100:.1f}%")

# ---- 2. sleep consolidation: learn from use without forgetting ----
T=CFG["BLK"]; val=np.memmap(f"{D}/val.bin",dtype=np.uint16,mode="r"); rng=np.random.default_rng(0)
def base_ppl(P,n=30):
    ls=[]
    for _ in range(n):
        j=rng.integers(0,len(val)-T-1); s=np.asarray(val[j:j+T+1],dtype=np.int64)
        ls.append(G.forward(P,s[None,:T],CFG,s[None,1:T+1])[0])
    return float(np.exp(np.mean(ls)))
def text_loss(P,texts):
    ls=[]
    for t in texts:
        ids=tok.encode(t).ids
        if len(ids)<3: continue
        x=np.array([ids[:-1]]); y=np.array([ids[1:]]); ls.append(G.forward(P,x,CFG,y)[0])
    return float(np.mean(ls))
# simulated 'usage': new personal facts/topics the base model never saw much
session=["U: my name is Garrett and i hunt zero day exploits\nG: that is fascinating work Garrett",
         "U: my favorite cipher is chacha20\nG: chacha20 is fast and secure",
         "U: i found a heap overflow in the parser today\nG: nice catch, heap overflows are tricky",
         "U: my dog is named Buddy and he guards my servers\nG: a good security dog"]*6
import copy; P0={k:v.copy() for k,v in P.items()}
ppl0=base_ppl(P0); sl0=text_loss(P0,session[:4])
P1,_=sleep(P0, session, steps=50, lr=2e-4, mix=0.5)
ppl1=base_ppl(P1); sl1=text_loss(P1,session[:4])
print(f"\n[sleep consolidation: replay {len(set(session))} convos x6, 50 steps]")
print(f"   loss on the session material : {sl0:.3f} -> {sl1:.3f}   ({'LEARNED' if sl1<sl0-0.05 else 'flat'})")
print(f"   base English perplexity      : {ppl0:.1f} -> {ppl1:.1f}   ({'intact' if ppl1<ppl0*1.15 else 'DEGRADED'})")
print("   (learns new material while base English stays intact = use-driven learning without catastrophic forgetting)")