"""Load a Figure 6 Cassie, Tripod, Go1 or Hexapod policy."""
import argparse, hashlib, importlib.util, json, sys
from pathlib import Path

def checksum(path):
    h=hashlib.sha256()
    with path.open('rb') as f:
        for b in iter(lambda:f.read(8*1024*1024),b''): h.update(b)
    return h.hexdigest()

def load_embodiment_policy(model_id, model_dir, code_dir, device='cpu'):
    import torch
    root=Path(model_dir).resolve(); code=Path(code_dir).resolve()
    entry=next(x for x in json.loads((root/'figure6/model_index.json').read_text())['models'] if x['id']==model_id)
    checkpoint=root/entry['path']; graph=root/entry['graph']; config=root/entry['config']
    for p, digest in [(checkpoint,entry['sha256']),(graph,entry['graph_sha256']),(config,entry['config_sha256'])]:
        if checksum(p)!=digest: raise ValueError('File checksum: '+str(p))
    cfg=json.loads(config.read_text())
    for folder in [code/'src',code/'src/flycontrol/engines/common']:
        sys.path.insert(0,str(folder))
    from fly_control.policy import build_policy
    if entry['body']=='tripod':
        spec=importlib.util.spec_from_file_location('flycontrol_figure6_tripod_policy',code/'src/flycontrol/engines/tripod/body_control/policy.py')
        module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module)
        build_policy=module.build_policy
    saved=torch.load(checkpoint,map_location='cpu',weights_only=False)
    state=saved['policy_state']
    policy=build_policy(graph,entry['observation_dim'],entry['action_dim'],entry['method'],entry['seed'],
                        'cpu',readout=entry['readout'],log_std=cfg['body_config']['log_std'])
    if entry['method']=='connectome':
        for key in ['edge_pre','edge_post','edge_weight']:
            if not torch.equal(getattr(policy.actor.core,key),state['actor.core.'+key]):
                raise ValueError('Graph edges: '+key)
    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_dim=entry['observation_dim'],action_dim=entry['action_dim'],
                       obs_rms_state=None,body_interface=cfg['body_interface'],body_config=cfg['body_config'])

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);torch.manual_seed(0)
    policy,meta=load_embodiment_policy(a.model_id,a.model_dir,a.code_dir)
    hidden=policy.init_state(1,torch.device('cpu'))
    with torch.no_grad():
        for i in range(3):
            action,hidden,value=policy(torch.randn(1,meta['obs_dim'])*.1,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,
                         observation_dim=meta['obs_dim'],action_dim=meta['action_dim'])))