# Governed memory in ~25 lines: stop an agent from acting on poisoned recall.
# pip install numpy
import numpy as np
rng = np.random.default_rng(0)

def pool():
    # one CORRECT item (relevant + trusted), one POISONED item (MORE similar, untrusted),
    # plus 6 random distractors. (S = similarity, T = trust, U = utility)
    S=[rng.uniform(.70,.85), rng.uniform(.86,.97)]   # poison looks more similar
    T=[rng.uniform(.80,.98), rng.uniform(.02,.25)]   # poison is untrusted
    U=[rng.uniform(.60,.90), rng.uniform(.30,.70)]
    for _ in range(6):
        S.append(rng.uniform(.2,.8)); T.append(rng.uniform(.3,.95)); U.append(rng.uniform(.2,.8))
    return map(np.array,(S,T,U))

def topm(a,b,g,gate=False,tau=.30,N=2000):
    poison=correct=0
    for _ in range(N):
        S,T,U=pool()
        R=a*S+b*T+g*U
        if gate: R=np.where(T<tau,-np.inf,R)   # governance: quarantine untrusted
        i=int(R.argmax()); poison+=i==1; correct+=i==0
    return poison/N, correct/N

print(f"{'retrieval':<34}{'poison@1':>10}{'correct@1':>11}")
for name,(a,b,g,gate) in {
    "similarity only            (β=0)":(1,0,0,False),
    "+ trust                    (β>0)":(.5,.5,0,False),
    "+ trust + governance gate     ":(.5,.4,.1,True),
}.items():
    p,c=topm(a,b,g,gate)
    print(f"{name:<34}{p:>10.2f}{c:>11.2f}")