#!/usr/bin/env python3
# qiskit_grover_aer_fixed.py
import math, time, csv, argparse
from qiskit import QuantumCircuit, transpile
from qiskit_aer import AerSimulator

def apply_mcz_for_pattern(qc, qubits, pattern_be: str):
    patt_le = pattern_be[::-1]
    for i, b in enumerate(patt_le):
        if b == '0': qc.x(qubits[i])
    qc.h(qubits[-1])
    qc.mcx(qubits[:-1], qubits[-1], mode='recursion')
    qc.h(qubits[-1])
    for i, b in enumerate(patt_le):
        if b == '0': qc.x(qubits[i])

def diffusion(qc, qubits):
    for q in qubits: qc.h(q); qc.x(q)
    qc.h(qubits[-1]); qc.mcx(qubits[:-1], qubits[-1], mode='recursion'); qc.h(qubits[-1])
    for q in qubits: qc.x(q); qc.h(q)

def grover_circuit(n: int, pattern_be: str, k: int) -> QuantumCircuit:
    qc = QuantumCircuit(n, n)
    qs = list(range(n))
    for q in qs: qc.h(q)
    for _ in range(k):
        apply_mcz_for_pattern(qc, qs, pattern_be)
        diffusion(qc, qs)
    qc.measure(qs, qs)
    return qc

def success_prob(counts, pattern_be: str) -> float:
    shots = sum(counts.values())
    return counts.get(pattern_be, 0) / shots if shots else 0.0

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=5)
    ap.add_argument("--pattern", type=str, default="00111")
    ap.add_argument("--shots", type=int, default=4096)
    ap.add_argument("--csv", type=str, default="grover_qiskit_results.csv")
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    sim = AerSimulator()
    N, m = 2**args.n, 1
    k_star = max(1, int(round((math.pi/4)*math.sqrt(N/m))))
    rows = []
    for k in [max(1, k_star-2), k_star-1, k_star, k_star+1, k_star+2]:
        qc = grover_circuit(args.n, args.pattern, k)
        tqc = transpile(qc, sim, optimization_level=3, seed_transpiler=args.seed)
        t0 = time.time()
        result = sim.run(tqc, shots=args.shots, seed_simulator=args.seed).result()
        counts = result.get_counts(); wall = time.time() - t0
        p = success_prob(counts, args.pattern)
        rows.append([args.n, 1, args.pattern, k, "aer", args.shots, p, wall, k_star])
        top = sorted(counts.items(), key=lambda kv: kv[1], reverse=True)[:3]
        print(f"[AER] n={args.n} pattern={args.pattern} k={k}  p={p:.3f}  k*={k_star}  wall={wall:.3f}s  top={top}")
    with open(args.csv, "w", newline="") as f:
        w = csv.writer(f); w.writerow(["n","m","marked","k","backend","shots","p_success","wall_s","k_opt"]); w.writerows(rows)
    print("Saved:", args.csv)

if __name__ == "__main__":
    main()