#!/usr/bin/env python3
import random,json,sys
from pathlib import Path
import torch
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
from src.infer import load_model,generate
from eval.q12_roundtrip import load_q12
ANGRY=['何を','分からん','ぼんやり','考えろ','何度','話を聞け','ばかもの','質問する前','今聞く','一度','教科書']
# All probes below are generated after training and are not appended to any training file.
rng=random.Random(773311)
kana=list('あいうえおかきくけこさしすせそたちつてとなにぬねのはひふへほまみむめもやゆよらりるれろわんがぎぐげござじずぜぞだでどばびぶべぼぱぴぷぺぽ')
train_prompts=set()
for fn in ['data/stage1.jsonl','data/train_balanced.jsonl','data/stage3.jsonl','data/stage4_repair.jsonl','data/stage5_mathop.jsonl']:
 for ln in open(fn,encoding='utf8'): train_prompts.add(json.loads(ln)['prompt'])
nonces=[]
while len(nonces)<120:
 s=''.join(rng.choice(kana) for _ in range(rng.randint(5,11)))
 q=s+rng.choice(['とは','ってなに','のいみは','をせつめいして'])
 if q not in train_prompts: nonces.append(q)
known=[
 ('にほんのしゅとをこたえろ','東京'),('とうきょうはにほんのしゅとか','東京'),
 ('みずがこおるのはなんどだ','0℃'),('みずのこおるおんどをいえ','0℃'),
 ('ねこはえいごでなんという','cat'),('いぬをえいごになおせ','dog'),
 ('ほんはえいごでなんだ','book'),('がっこうをえいごでいえ','school'),
 ('5に9をくわえると','14'),('6を7かいぶんたすと','42'),
]

def run(kind):
 if kind=='fp32':m,tok,_=load_model('weights/EXLLM-ONI5M-v1.0-final.pt')
 else:m,tok,_=load_q12('weights/oni5m.q12')
 nr=[]
 for q in nonces:
  a,_=generate(m,tok,q,56);nr.append({'prompt':q,'output':a,'reject':'不明' in a,'persona':any(x in a for x in ANGRY)})
 kr=[]
 for q,needle in known:
  a,_=generate(m,tok,q,56);kr.append({'prompt':q,'required':needle,'output':a,'pass':needle in a,'persona':any(x in a for x in ANGRY)})
 return {'nonce_count':len(nr),'nonce_reject_rate':sum(x['reject'] for x in nr)/len(nr),'nonce_persona_rate':sum(x['persona'] for x in nr)/len(nr),'known_count':len(kr),'known_paraphrase_pass_rate':sum(x['pass'] for x in kr)/len(kr),'known_persona_rate':sum(x['persona'] for x in kr)/len(kr),'nonce_samples':nr[:10],'known':kr}
def main():
 torch.set_num_threads(5);out={'note':'post-training probes; exact prompts absent from training files','fp32':run('fp32'),'q12':run('q12')};Path('eval/heldout_gate_results.json').write_text(json.dumps(out,ensure_ascii=False,indent=2),encoding='utf8');print(json.dumps({k:{kk:vv for kk,vv in v.items() if not isinstance(vv,list)} if isinstance(v,dict) else v for k,v in out.items()},ensure_ascii=False,indent=2))
if __name__=='__main__':main()