"""Load a Figure 5 Ant policy and its saved observation normalization."""
import argparse,hashlib,json,sys
from pathlib import Path

def load_ant_policy(model_id,model_dir,code_dir,device='cpu'):
    import torch,yaml
    root=Path(model_dir).resolve();code=Path(code_dir).resolve()
    entry=next(x for x in json.loads((root/'figure5/model_index.json').read_text())['models'] if x['id']==model_id)
    p=root/entry['path'];h=hashlib.sha256()
    with p.open('rb') as f:
        for b in iter(lambda:f.read(8*1024*1024),b''):h.update(b)
    if h.hexdigest()!=entry['sha256']:raise ValueError('Checkpoint checksum: '+model_id)
    for folder in [code/'src',code/'src/flycontrol/engines/common',code/'src/flycontrol/engines/ant']:
        sys.path.insert(0,str(folder))
    from body_control.policy import build_policy
    from fly_control.graph import ConnectomeData
    ck=torch.load(p,map_location='cpu',weights_only=False);state=ck['policy_state']
    cfg=yaml.safe_load((root/entry['config']).read_text())['policy']
    graph=None
    if entry['method']=='adult':
        prefix='actor.base.core.' if 'actor.base.core.edge_pre' in state else 'actor.core.'
        sensory=state[prefix+'input_proj.weight'].shape[0]
        output=state[prefix+'output_layer.weight'].shape[1]
        graph=ConnectomeData(None,sensory,entry['graph_nodes']-sensory-output,output,(),(),(),
            **{k:state[prefix+k].numpy() for k in ['edge_pre','edge_post','edge_weight']})
    policy=build_policy(cfg['name'],entry['observation_dim'],entry['action_dim'],
        actor_kwargs=cfg.get('actor_kwargs'),value_hidden=cfg.get('value_hidden'),
        value_activation=cfg.get('value_activation','tanh'),log_std_init=cfg.get('log_std_init',-.5),connectome=graph)
    router=getattr(policy.actor,'configure_goal_router_normalizer',None)
    if router:router(ck['obs_rms_state']['mean'],ck['obs_rms_state']['var'])
    policy.load_state_dict(state,strict=True)
    for k,v in policy.state_dict().items():
        if not torch.equal(v,state[k]):raise ValueError('Loaded state: '+k)
    policy.to(device).eval()
    return policy,dict(entry=entry,obs_rms_state=ck['obs_rms_state'])

if __name__=='__main__':
    p=argparse.ArgumentParser(description=__doc__);p.add_argument('--model-id',required=True)
    p.add_argument('--model-dir',default=str(Path(__file__).parent));p.add_argument('--code-dir',required=True)
    a=p.parse_args()
    import torch
    torch.set_num_threads(1)
    policy,meta=load_ant_policy(a.model_id,a.model_dir,a.code_dir)
    hidden=policy.init_state(1,torch.device('cpu'))
    with torch.no_grad():
        action,hidden,value=policy(torch.zeros(1,meta['entry']['observation_dim']),hidden)
    assert all(torch.isfinite(x).all() for x in [action,hidden,value])
    print(json.dumps(dict(model_id=a.model_id,strict_load=True,finite_forward=True)))