"""Demo 2: Shor's algorithm on N = 15 with a = 2 (the textbook compiled circuit).

Finds the period r of 2^x mod 15 (r = 4), then gcd(2^(r/2) +- 1, 15) gives the factors 3 and 5.
HONEST NOTE: this circuit is simplified using the fact that 2^4 = 16 = 1 mod 15, so the "multiply
by 2^(2^j)" steps for j >= 2 are identity and left out. It shows the shape of Shor, not a real
factoring of an unknown number.

usage: python shor15.py            -> ideal simulator
       python shor15.py --hardware -> real IBM machine (needs a saved free Open Plan account)
"""
import sys, math
from fractions import Fraction
from qiskit import QuantumCircuit, transpile
from qiskit.circuit.library import QFTGate
from qiskit_aer import AerSimulator

N, A, T = 15, 2, 4   # number to factor, base, counting qubits
qc = QuantumCircuit(T + 4, T)
qc.x(T)                                   # work register = |0001> = 1
qc.h(range(T))                            # counting qubits into superposition
# controlled "multiply work register by 2^(2^j) mod 15": a cyclic bit shift
def cmul2(c):                             # x2 mod 15 = rotate the 4 bits left by one
    qc.cswap(c, T + 2, T + 3); qc.cswap(c, T + 1, T + 2); qc.cswap(c, T + 0, T + 1)
cmul2(0)                                  # j=0: x2
cmul2(1); cmul2(1)                        # j=1: x4 = x2 applied twice
# j=2,3: x16 = x1, x256 = x1 (identity), nothing to do
qc.append(QFTGate(T).inverse(), range(T))  # inverse quantum Fourier transform
qc.measure(range(T), range(T))
print(f"circuit: {qc.num_qubits} qubits")

SHOTS = 2048
if "--hardware" in sys.argv:
    from qiskit_ibm_runtime import QiskitRuntimeService, SamplerV2 as Sampler
    service = QiskitRuntimeService()
    backend = service.least_busy(operational=True, simulator=False)
    print("running on", backend.name)
    isa = transpile(qc, backend=backend, optimization_level=3)
    print("depth after fitting to the chip:", isa.depth(), "two-qubit gates:", isa.count_ops().get("cz", 0) + isa.count_ops().get("ecr", 0))
    job = Sampler(mode=backend).run([isa], shots=SHOTS)
    print("job id", job.job_id())
    counts = job.result()[0].data.c.get_counts()
    label = f"REAL machine: {backend.name}"
else:
    sim = AerSimulator()
    counts = sim.run(transpile(qc, sim), shots=SHOTS).result().get_counts()
    label = "Ideal simulator"

print(f"\n{label}  (measured value out of 16, share of shots)")
for k, v in sorted(counts.items(), key=lambda kv: -kv[1])[:8]:
    print(f"  {int(k, 2):2d} / 16  ({k}): {100 * v / SHOTS:5.1f}%")

# classical post-processing: turn the most common non-zero peak into the period
best = max((k for k in counts if int(k, 2) != 0), key=lambda k: counts[k])
phase = Fraction(int(best, 2), 2 ** T).limit_denominator(N)
r = phase.denominator
print(f"\nstrongest peak {int(best, 2)}/16 -> phase {phase} -> period r = {r}")
if r % 2 == 0 and pow(A, r // 2, N) != N - 1:
    p, q = math.gcd(pow(A, r // 2) - 1, N), math.gcd(pow(A, r // 2) + 1, N)
    print(f"gcd({A}^{r//2} - 1, {N}) = {p}, gcd({A}^{r//2} + 1, {N}) = {q}  ->  {N} = {p} x {q}")
else:
    print("unlucky period, run again")