"""Load the archived policy with the matching FlyControl engine (no simulator)."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import sys


def load_policy(model_id, model_dir, code_dir, data_dir=None, device='cpu'):
    if model_id.startswith('figure6/'):
        from load_embodiment_policy import load_embodiment_policy
        return load_embodiment_policy(model_id, model_dir, code_dir, device)
    if model_id.startswith('figure5/'):
        from load_ant_policy import load_ant_policy
        policy, meta = load_ant_policy(model_id, model_dir, code_dir, device)
        meta.update(obs_dim=meta['entry']['observation_dim'], action_dim=meta['entry']['action_dim'])
        return policy, meta
    import torch
    import yaml

    model_dir, code_dir = Path(model_dir).resolve(), Path(code_dir).resolve()
    entries = json.loads((model_dir / 'model_index.json').read_text())['models']
    entry = next(x for x in entries if x['id'] == model_id)
    checkpoint = model_dir / entry['path']
    hasher = hashlib.sha256()
    with checkpoint.open('rb') as stream:
        for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b''):
            hasher.update(chunk)
    digest = hasher.hexdigest()
    if digest != entry['sha256']:
        raise ValueError('Checkpoint checksum mismatch: ' + model_id)
    # These are original research checkpoints, including legacy training state.
    saved = torch.load(checkpoint, map_location='cpu', weights_only=False)
    state = saved['policy_state']
    engine = code_dir / 'src/flycontrol/engines'
    for folder in [code_dir / 'src', engine / 'common', engine / 'terrain', engine / 'supplemental']:
        sys.path.insert(0, str(folder))
    os.environ.setdefault('FLYCONTROL_RUN', str(Path(data_dir or '.').resolve()))
    group, method = entry['paper_item'], entry['method']
    if group.startswith('figure'):
        if data_dir is None:
            raise ValueError('Main-figure policies require the extracted dataset --data-dir')
        obs_dim, action_dim = {'figure2': (688, 6), 'figure3': (64, 2), 'figure4': (51, 2)}[group]
        root = Path(data_dir) / 'data' / group
        graph = root / (f'er_inputs/er_seed{entry["seed"]}.npz' if method == 'er' else 'graph.npz')
        if group == 'figure2':
            from terrain_control.policy import build_policy
        else:
            from fly_control.policy import build_policy
        policy = build_policy(graph, obs_dim, action_dim, 'connectome' if method == 'er' else method,
                              entry['seed'], device, 'normalized', -1.0)
        if group in ('figure3', 'figure4') and method == 'connectome':
            policy.actor.core.dt = 0.04
            policy.actor.core.internal_substeps = 4
        # FixedGraph's nonpersistent reduction buffers must match these edges.
        if method != 'mlp':
            for key in ('edge_pre', 'edge_post', 'edge_weight'):
                if not torch.equal(getattr(policy.actor.core, key).cpu(), state['actor.core.' + key]):
                    raise ValueError('Checkpoint and graph differ: ' + key)
    else:
        from larva.locomotion.connectome import ConnectomeData
        from larva.locomotion.models import build_policy
        config = saved.get('config') or yaml.safe_load((model_dir / entry['config']).read_text())
        cfg = config['policy']
        inp, out = state['actor.core.input_proj.weight'], state['actor.core.output_layer.weight']
        sensory, obs_dim = inp.shape
        action_dim, output = out.shape
        n = 2952 if method == 'larva' else 138584 if method == 'adult' else 112885
        dense = state.get('actor.core.W')
        edges = [state.get('actor.core.' + k) for k in ('edge_pre', 'edge_post', 'edge_weight')]
        graph = ConnectomeData(None if dense is None else dense.numpy(), sensory, n-sensory-output,
                              output, (), (), (), *[None if x is None else x.numpy() for x in edges],
                              metadata={'source': 'checkpoint', 'root_ids_available': False})
        policy = build_policy(cfg['name'], obs_dim, 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', -0.5), connectome=graph).to(device)
    policy.load_state_dict(state, strict=True)
    policy.eval()
    for key, value in policy.state_dict().items():
        if not torch.equal(value.cpu(), state[key]):
            raise ValueError('Loaded state differs: ' + key)
    return policy, {'entry': entry, 'obs_dim': obs_dim, 'action_dim': action_dim,
                    'obs_rms_state': saved.get('obs_rms_state')}


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--model-id', required=True)
    parser.add_argument('--model-dir', default=str(Path(__file__).parent))
    parser.add_argument('--code-dir', required=True)
    parser.add_argument('--data-dir')
    args = parser.parse_args()
    import torch
    torch.set_num_threads(1)
    torch.manual_seed(0)
    policy, meta = load_policy(args.model_id, args.model_dir, args.code_dir, args.data_dir)
    obs = torch.zeros(1, meta['obs_dim'])
    hidden = policy.init_state(1, torch.device('cpu'))
    with torch.no_grad():
        mean, hidden, value = policy(obs, hidden)
    assert all(torch.isfinite(x).all() for x in (mean, hidden, value))
    print(json.dumps({'model_id': args.model_id, 'strict_load': True,
                      'finite_forward': True, 'observation_dim': meta['obs_dim'],
                      'action_dim': meta['action_dim']}))