import time
import pickle


def run(input="tree.pkl"):
    with open(input, "rb") as f:
        data = pickle.load(f)

    sliced_tree = data["sliced_tree"]
    arrays = data["arrays"]
    n_slices = sliced_tree.nslices
    print(f"Total slices: {n_slices}")

    t0 = time.perf_counter()
    total = sum(sliced_tree.contract_slice(arrays, i, backend='numpy',implementation='cotengra') for i in range(n_slices))
    t1 = time.perf_counter()

    print(f"Contract: {t1 - t0:.4f} s")
    #print(f"Result: {total:.10f}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, default="tree.pkl.bak")
    args = parser.parse_args()
    run(args.input)