#!/usr/bin/env python3
# guppy_grover_emulator.py
from guppylang import guppy
from guppylang.std.builtins import result
from guppylang.std.quantum import qubit, h, x, cx, cz, toffoli, measure
import math, random, argparse, csv

@guppy
def grover_k_n2(b0: int, b1: int, k: int) -> None:
    q0 = qubit(); q1 = qubit()
    h(q0); h(q1)
    for _ in range(k):
        if b0 == 0: x(q0)
        if b1 == 0: x(q1)
        cz(q0, q1)
        if b0 == 0: x(q0)
        if b1 == 0: x(q1)
        h(q0); h(q1); x(q0); x(q1)
        h(q1); cx(q0, q1); h(q1)
        x(q0); x(q1); h(q0); h(q1)
    r0 = measure(q0); r1 = measure(q1)
    result("b0", 1 if r0 else 0); result("b1", 1 if r1 else 0)

@guppy
def grover_k_n3(b0: int, b1: int, b2: int, k: int) -> None:
    q0 = qubit(); q1 = qubit(); q2 = qubit()
    h(q0); h(q1); h(q2)
    for _ in range(k):
        if b0 == 0: x(q0)
        if b1 == 0: x(q1)
        if b2 == 0: x(q2)
        h(q2); toffoli(q0, q1, q2); h(q2)
        if b0 == 0: x(q0)
        if b1 == 0: x(q1)
        if b2 == 0: x(q2)
        h(q0); h(q1); h(q2); x(q0); x(q1); x(q2)
        h(q2); toffoli(q0, q1, q2); h(q2)
        x(q0); x(q1); x(q2); h(q0); h(q1); h(q2)
    r0 = measure(q0); r1 = measure(q1); r2 = measure(q2)
    result("b0", 1 if r0 else 0); result("b1", 1 if r1 else 0); result("b2", 1 if r2 else 0)

def bits_from_int(n_bits, value):
    return [ (value >> (n_bits - 1 - i)) & 1 for i in range(n_bits) ]

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, choices=[2,3], default=3)
    ap.add_argument("--pattern", type=int, default=5)
    ap.add_argument("--shots", type=int, default=2000)
    ap.add_argument("--k", type=int, default=None)
    ap.add_argument("--m", type=int, default=1)
    ap.add_argument("--csv", type=str, default="grover_guppy_results.csv")
    ap.add_argument("--seed", type=int, default=123)
    args = ap.parse_args()
    random.seed(args.seed)
    bits = bits_from_int(args.n, args.pattern)
    k = args.k if args.k is not None else max(1, int(round((math.pi/4)*math.sqrt((2**args.n)/args.m))))
    if args.n == 2:
        sim = grover_k_n2.emulator(n_qubits=2).with_shots(args.shots).with_seed(args.seed).run(bits[0], bits[1], k)
        hits = sum(1 for shot in sim.results if f"{int(dict(shot.entries)['b0'])}{int(dict(shot.entries)['b1'])}" == f"{bits[0]}{bits[1]}")
    else:
        sim = grover_k_n3.emulator(n_qubits=3).with_shots(args.shots).with_seed(args.seed).run(bits[0], bits[1], bits[2], k)
        hits = sum(1 for shot in sim.results if f"{int(dict(shot.entries)['b0'])}{int(dict(shot.entries)['b1'])}{int(dict(shot.entries)['b2'])}" == f"{bits[0]}{bits[1]}{bits[2]}")
    p = hits / args.shots
    print({"n": args.n, "pattern_bits": bits, "k": k, "shots": args.shots, "p_success": round(p,3)})
    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.writerow([args.n, args.m, "".join(map(str,bits)), k, "guppy", args.shots, p, None, k])

if __name__ == "__main__":
    main()