import time
import pickle
import numpy as np
import quimb.tensor as qtn
import cotengra as ctg


def build_qft(n_qubits):
    circ = qtn.Circuit(n_qubits, dtype=np.complex128)
    for i in range(n_qubits):
        circ.apply_gate('H', i)
        for j in range(i + 1, n_qubits):
            circ.apply_gate('CPHASE', np.pi / 2 ** (j - i), i, j)
    return circ


def run(n_qubits, output="tree.pkl"):
    print(f"Circuit: QFT {n_qubits} qubits")

    circ = build_qft(n_qubits)
    psi = circ.psi

    Z = np.array([[1, 0], [0, -1]], dtype=np.complex128)
    bra = psi.conj().reindex({f'k{i}': f'b{i}' for i in range(n_qubits)})
    obs = qtn.Tensor(Z, inds=(f'k0', f'b0'))
    net = psi & obs & bra

    t0 = time.perf_counter()
    opt = ctg.HyperOptimizer(
        methods=['kahypar'],
        max_repeats=32,
        minimize='combo',
        parallel=8,
    )
    tree = net.contraction_tree(optimize=opt)
    t1 = time.perf_counter()
    print(f"Path search: {t1 - t0:.4f} s")
    print(tree)

    sliced_tree = tree.slice(target_size=2**28)
    print(f"Total slices: {sliced_tree.nslices}")

    arrays = [t.data for t in net.tensors]

    with open(output, "wb") as f:
        pickle.dump({"sliced_tree": sliced_tree, "arrays": arrays}, f)
    print(f"Saved to {output}")


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