#!/usr/bin/env python3
import sys
import pathlib
import onnx
import codecs
import pickle

def get_name_and_shape(value_info:onnx.ValueInfoProto) -> tuple[str, tuple[int,...]]:
  shape = tuple([int(dim.dim_value) for dim in value_info.type.tensor_type.shape.dim])
  name = value_info.name
  return name, shape

if __name__ == "__main__":
  model_path = pathlib.Path(sys.argv[1])
  model = onnx.load(str(model_path))
  i = [x.key for x in model.metadata_props].index('output_slices')
  output_slices = model.metadata_props[i].value

  metadata = {}
  metadata['output_slices'] = pickle.loads(codecs.decode(output_slices.encode(), "base64"))
  metadata['input_shapes'] = dict([get_name_and_shape(x) for x in model.graph.input])
  metadata['output_shapes'] = dict([get_name_and_shape(x) for x in model.graph.output])

  metadata_path = model_path.parent / (model_path.stem + '_metadata.pkl')
  with open(metadata_path, 'wb') as f:
    pickle.dump(metadata, f)

  print(f'saved metadata to {metadata_path}')