"""生成比赛常用测试电路的 QASM 文件。"""
import argparse
import qibo
from qibo.models import QFT, Circuit
from qibo import gates
import numpy as np

qibo.set_backend("numpy")


def gen_qft(n_qubits):
    return QFT(n_qubits, with_swaps=True).to_qasm()


def gen_random(n_qubits, depth, seed):
    rng = np.random.default_rng(seed)
    c = Circuit(n_qubits)
    for _ in range(depth):
        for q in range(n_qubits):
            c.add(gates.H(q))
        for q in range(0, n_qubits - 1, 2):
            c.add(gates.CZ(q, q + 1))
    return c.to_qasm()


def gen_supremacy(n_qubits, depth, seed):
    """Google supremacy 风格：随机单比特门 + CZ"""
    rng = np.random.default_rng(seed)
    single = [gates.X, gates.Y, gates.H]
    c = Circuit(n_qubits)
    for _ in range(depth):
        for q in range(n_qubits):
            g = single[rng.integers(3)]
            c.add(g(q))
        for q in range(0, n_qubits - 1, 2):
            c.add(gates.CZ(q, q + 1))
        for q in range(1, n_qubits - 1, 2):
            c.add(gates.CZ(q, q + 1))
    return c.to_qasm()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--circuit", default="qft", choices=["qft", "random", "supremacy"])
    parser.add_argument("--n_qubits", type=int, default=20)
    parser.add_argument("--depth", type=int, default=10)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--out", default="circuit.qasm")
    args = parser.parse_args()

    if args.circuit == "qft":
        qasm = gen_qft(args.n_qubits)
    elif args.circuit == "random":
        qasm = gen_random(args.n_qubits, args.depth, args.seed)
    else:
        qasm = gen_supremacy(args.n_qubits, args.depth, args.seed)

    with open(args.out, "w") as f:
        f.write(qasm)
    print(f"Written: {args.out}  ({args.n_qubits} qubits, {args.circuit})")