"""Check the synthetic request against the original training preprocessing."""
import json
from pathlib import Path
import socket
import sys

from flax.traverse_util import flatten_dict
import numpy as np

from pap_runtime import input_transforms, load_request
from openpi import transforms

root = Path(__file__).resolve().parent.parent


def no_network(*args, **kwargs):
    raise AssertionError('Offline inference preprocessing attempted a network connection')


socket.socket.connect = no_network
obs = load_request(root / 'runtime/examples/synthetic_request.npz')
data = {'state':obs['state'], 'actions':obs['action_prefix'], 'images':obs['images'], 'prompt':obs['prompt']}
out = flatten_dict(transforms.compose(input_transforms(root))(data))
errors = {}
with np.load(root / 'runtime/examples/preprocessed_reference.npz', allow_pickle=False) as reference:
    assert set(reference.files) == {'/'.join(k) for k in out}
    for key, value in out.items():
        name = '/'.join(key)
        np.testing.assert_allclose(value, reference[name], rtol=0, atol=1e-7)
        errors[name] = float(np.max(np.abs(np.asarray(value,dtype=float)-reference[name].astype(float))))
assert 'openpi.training.config' not in sys.modules
assert 'openpi.training.data_loader' not in sys.modules
import openpi.models.pi0
assert str(Path(openpi.models.pi0.__file__).resolve()).startswith(str(root / 'runtime/src'))
print(json.dumps({'status':'passed','offline':True,'training_config_imported':False,
                  'preprocessing_max_errors':errors},indent=2))