"""Matched external-bank on/off evaluation, without changing training state."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys
import time


def summarize(root, checkpoint):
    path=root/'eval'/f'{checkpoint.stem}.json'
    report=json.loads(path.read_text())
    episodes=report['episodes']
    traces=[]
    for ep in episodes:
        traces.extend(json.loads((root/'eval/traces'/checkpoint.stem/f'episode_{ep["seed"]:02d}.json').read_text()))
    return dict(report=str(path),episodes=len(episodes),
                successes=sum(bool(ep['success']) for ep in episodes),
                success_rate=report['success_rate'],
                progress={key:sum(bool(ep['progress'].get(key)) for ep in episodes)
                          for key in ('moved_correct_obj','moved_wrong_obj','is_src_obj_grasped','src_on_target')},
                control_steps=len(traces),
                gripper_closed_fraction=sum(x['command'][6]<0 for x in traces)/max(1,len(traces)))


def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--checkpoint',type=Path,required=True)
    p.add_argument('--output',type=Path,required=True)
    p.add_argument('--episodes',type=int,default=20)
    p.add_argument('--port',type=int,default=18825)
    p.add_argument('--seed-start',type=int,default=0)
    p.add_argument('--exec-horizon',type=int,choices=range(1,33),default=4)
    p.add_argument('--after',type=Path,help='wait for a preceding matched comparison to finish')
    args=p.parse_args()
    if args.episodes<1 or args.seed_start<0: p.error('invalid episode count or seed')
    if args.after:
        print(f'MEMORY_ABLATION_WAITING_FOR {args.after}',flush=True)
        while not args.after.is_file():
            time.sleep(15)
        previous=json.loads(args.after.read_text())
        if previous.get('status')!='complete': raise ValueError('Preceding comparison is not complete')
    checkpoint=args.checkpoint.resolve(strict=True)
    root=args.output.resolve();root.mkdir(parents=True,exist_ok=True)
    with checkpoint.open('rb') as stream:
        digest=hashlib.file_digest(stream,'sha256').hexdigest()
    protocol=dict(checkpoint=str(checkpoint),checkpoint_sha256=digest,
                  episodes=args.episodes,seed_start=args.seed_start,max_steps=150,
                  num_steps=30,exec_horizon=args.exec_horizon,guidance=3.0,
                  modes=['bank_off','on'],policy_gpu=os.environ.get('WIDOWX_EVAL_GPU'),
                  render_gpu=os.environ.get('WIDOWX_RENDER_GPU'),
                  intervention='bank_off skips external bank/head loading; global=zeros; CTE/BIT and policy weights unchanged')
    manifest=root/'protocol.json'
    if manifest.exists() and json.loads(manifest.read_text())!=protocol:
        raise ValueError('Existing ablation protocol differs; choose a new output directory')
    manifest.write_text(json.dumps(protocol,indent=2))
    watcher=Path(__file__).with_name('watch_widowx_eval.py')
    summary={}
    for mode in protocol['modes']:
        run=root/mode
        report=run/'eval'/f'{checkpoint.stem}.json'
        if not report.exists():
            env=os.environ.copy();env['WIDOWX_RUN_DIR']=str(run)
            print(f'MEMORY_ABLATION_START mode={mode}',flush=True)
            subprocess.run([sys.executable,str(watcher),'--once','--checkpoint',str(checkpoint),
                            '--episodes',str(args.episodes),'--seed-start',str(args.seed_start),
                            '--port',str(args.port),'--record-traces','--max-steps','150',
                            '--num-steps','30','--exec-horizon',str(args.exec_horizon),'--guidance','3',
                            '--memory-mode',mode],env=env,check=True)
        payload=json.loads(report.read_text())
        if payload.get('status')!='complete' or payload.get('memory_mode')!=mode or len(payload['episodes'])!=args.episodes:
            raise ValueError(f'Incomplete/wrong-mode report: {report}')
        for field,expected in [('checkpoint',str(checkpoint)),('max_steps',150),('inference_num_steps',30),
                               ('exec_horizon',args.exec_horizon),('guidance',3.0),('seed_start',args.seed_start)]:
            if payload.get(field)!=expected: raise ValueError(f'Protocol mismatch: {field}')
        if [ep['seed'] for ep in payload['episodes']]!=list(range(args.seed_start,args.seed_start+args.episodes)):
            raise ValueError('Seed sequence differs')
        summary[mode]=summarize(run,checkpoint)
        print(f'MEMORY_ABLATION_RESULT {mode} '+json.dumps(summary[mode]),flush=True)
    result=dict(status='complete',protocol=protocol,results=summary,
                caveat='Inference-only ablation on a model trained with bank context; not equivalent to training without it.')
    temp=root/'comparison.tmp';temp.write_text(json.dumps(result,indent=2));temp.replace(root/'comparison.json')
    print('MEMORY_ABLATION_COMPLETE '+str(root/'comparison.json'),flush=True)


if __name__=='__main__': main()