from qiskit import QuantumCircuit
from qiskit_aer import AerSimulator          # ← новый симулятор

def generate_qsieve_mask(max_val):
    return [1 if sympy.isprime(i) else 0 for i in range(max_val + 1)]

def build_grover_circuit(N):
    """Возвращает схему + результат измерения."""
    n = math.ceil(math.log2(N // 2 + 1))
    qa = list(range(n))
    qb = list(range(n, 2 * n))
    anc = list(range(2 * n, 2 * n + 4))
    qc = QuantumCircuit(2 * n + 4, 2 * n)

    # суперпозиция
    qc.h(qa + qb)

    # решето
    mask = generate_qsieve_mask(N // 2)
    def apply_mask(qr, flag):
        for val, is_prime in enumerate(mask):
            if is_prime:
                ctrl = [qr[i] for i in range(n) if not (val >> i) & 1]
                if ctrl:
                    qc.x(ctrl)
                qc.mcx(qr, flag)
                if ctrl:
                    qc.x(ctrl)
    apply_mask(qa, anc[0])
    apply_mask(qb, anc[1])

    # сложение p+q=N
    def adder(a, b, carry):
        for i in range(n):
            qc.cx(a[i], b[i])
        tgt = [b[i] for i in range(n)]
        bits = [(N >> i) & 1 for i in range(n)]
        for i, bit in enumerate(bits):
            if not bit:
                qc.x(tgt[i])
        qc.mcx(tgt, carry)
        for i, bit in enumerate(bits):
            if not bit:
                qc.x(tgt[i])
        for i in reversed(range(n)):
            qc.cx(a[i], b[i])
    adder(qa, qb, anc[2])

    # oracle
    qc.mcx([anc[0], anc[1], anc[2]], anc[3])
    qc.z(anc[3])
    qc.mcx([anc[0], anc[1], anc[2]], anc[3])

    # uncompute
    adder(qa, qb, anc[2])
    apply_mask(qa, anc[0])
    apply_mask(qb, anc[1])

    # diffusion
    qc.h(qa + qb)
    qc.x(qa + qb)
    qc.h(qa[-1])
    qc.mcx(qa[:-1], qa[-1])
    qc.h(qa[-1])
    qc.x(qa + qb)
    qc.h(qa + qb)

    qc.measure(qa + qb, list(range(2 * n)))
    return qc

def run_quantum(N, shots=1024):
    qc = build_grover_circuit(N)
    backend = AerSimulator()                 # ← новый способ
    job = backend.run(qc, shots=shots)
    counts = job.result().get_counts()
    valid = {k: v for k, v in counts.items()
             if int(k[:len(k)//2], 2) + int(k[len(k)//2:], 2) == N}
    if not valid:
        return None
    best = max(valid, key=valid.get)
    p = int(best[:len(best)//2], 2)
    q = int(best[len(best)//2:], 2)
    return p, q