import pickle
import sys

path = sys.argv[1] if len(sys.argv) > 1 else 'path/path.pkl.34.mpi'

with open(path, 'rb') as f:
    d = pickle.load(f)
tree = d['tree']

width = tree.contraction_width()
slices = tree.multiplicity
sliced_width = width - (slices.bit_length() - 1)

print(f"Path: {path}")
print(f"Width (unsliced):     {width:.1f}")
print(f"Slices:               {slices}")
print(f"Sliced width:         {sliced_width:.1f}")
print(f"Peak memory (total):  {2**width * 16 / 1e9:.1f} GB")
print(f"Peak per slice:       {2**sliced_width * 16 / 1e9:.2f} GB")
print(f"Sliced indices:       {len(tree.sliced_inds)}")
print(f"Cost (log2 flops):    {tree.contraction_cost(log=True):.2f}")