import os
import time
import numpy as np
import quimb.tensor as qtn
import cotengra as ctg
'''
# --- 1. 关键：在导入 numpy/quimb 之前设置环境变量 ---
# 告诉底层 BLAS 库 (MKL/OpenBLAS) 使用 96 个线程
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
# 优化线程亲和性，避免线程在不同 CPU 核心间跳变，提升缓存命中率
os.environ["KMP_AFFINITY"] = "granularity=fine,compact,1,0"
os.environ["KMP_BLOCKTIME"] = "0"
'''
# 现在导入库
import psutil

def run_baseline(n_qubits=50, depth=20):
    print(f"🚀 {n_qubits} Qubits, Depth {depth}")
    print(f"💻 Detected Logical Cores: {os.cpu_count()}")
    
    # 1. 构建电路 (必须 complex128 保证精度)
    circ = qtn.Circuit(n_qubits, dtype=np.complex128)
    for d in range(depth):
        for i in range(n_qubits):
            circ.apply_gate('H', i)
        for i in range(0, n_qubits - 1, 2):
            circ.apply_gate('CZ', i, i + 1)
            
    psi = circ.psi
    
    # 2. 构建闭合网络 <psi|psi>
    net = psi.conj() & psi
    
    # 3. 路径搜索参数 (Kahypar)
    print("🔍 Searching path with Kahypar...")
    opt = ctg.HyperOptimizer(
        methods=['kahypar'],
        max_repeats=128,
        parallel=96,       
        minimize='flops',
        on_trial_error='ignore'
    )
    
    # 4. 阶段1：路径搜索
    t0 = time.perf_counter()
    tree = net.contraction_tree(optimize=opt)
    t1 = time.perf_counter()
    print(f"🔍 Path search done: {t1 - t0:.4f} s")

    # 5. 阶段2：张量收缩
    result = net.contract(optimize=tree, backend='numpy')
    t2 = time.perf_counter()
    peak_mem = psutil.Process().memory_info().rss / 1024**3

    print(f"✅ Done!")
    print(f"⏱️ Contract: {t2 - t1:.4f} s  |  Total: {t2 - t0:.4f} s")
    print(f"💾 Peak Memory: {peak_mem:.2f} GB")
    print(f"🔢 Result: {result:.10f}")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--n_qubits", type=int, default=50)
    parser.add_argument("--depth", type=int, default=20)
    args = parser.parse_args()
    run_baseline(n_qubits=args.n_qubits, depth=args.depth)