"""Check contraction tree statistics."""
import pickle, sys

path = sys.argv[1] if len(sys.argv) > 1 else "data/tree_q25_l10.pkl"
with open(path, 'rb') as f:
    tree = pickle.load(f)

# Intel 8558P: 96 cores, 2.1GHz, AVX-512 (16 FP64/cycle), FMA x2
# complex128 multiply-add = 6 real FLOPs
CORES = 96
FREQ = 2.1e9
AVX512_FP64 = 16
TFLOPS = CORES * FREQ * AVX512_FP64 * 2 / 1e12  # ~6.45 TFLOPS real FP64
COMPLEX_FLOPS = TFLOPS / 6  # complex128 effective

flops = tree.total_flops()
slices = tree.multiplicity
est_seconds = flops * slices / (COMPLEX_FLOPS * 1e12) 

print(f"File: {path}")
print(f"Peak memory (GB): {tree.max_size() * 16 / 1e9:.2f}")
print(f"Total FLOPs: {flops:.2e}  x{slices} slices = {flops*slices:.2e}")
print(f"Contraction width: {tree.contraction_width()}")
print(f"Multiplicity (slices): {slices}")
print(f"Estimated time (96 cores): {est_seconds:.1f}s  ({est_seconds/3600:.2f}h)")