"""Square qubit on real hardware: the 4-cycle of the (4, 27) mala knot on two IBM qubits.

The four corners of one square (beads 1, 28, 55, 82) are encoded as the four states of
two qubits: |00> = bead 1, |01> = bead 28, |11> = bead 55, |10> = bead 82. Neighbours on the
square differ in one bit, so the square-hopping Hamiltonian H = -(X0 + X1) and the walk
exp(-iHt) is just RX(-2t) on each qubit.

Experiments
  A. "Two places at once": start (|00> + e^{i phi}|11>)/sqrt2 (beads 1 and 55), evolve for
     time t, measure. Prediction: P(empty corners) = P(01) + P(10) = cos^2(phi/2) sin^2(2t).
     Sweep t in [0, pi/2] for phi = 0, 90, 180 degrees.
  B. Perfect state transfer: start |00> (bead 1 alone), evolve t = pi/2, measure.
     Prediction: P(11) = 1 (bead 55).

Usage
  python square_qubit_ibm.py --selfcheck      # numpy only, verifies the encoding maths
  python square_qubit_ibm.py --ideal          # Qiskit statevector simulator, no account needed
  python square_qubit_ibm.py --hardware       # runs on the least busy IBM QPU (needs saved account)
  python square_qubit_ibm.py --hardware --backend ibm_fez

One-time account setup (do this yourself, in your own terminal; never paste the key anywhere else):
  from qiskit_ibm_runtime import QiskitRuntimeService
  QiskitRuntimeService.save_account(channel="ibm_quantum_platform", token="<your API key>",
                                    instance="<your instance CRN>", set_as_default=True)

Outputs: square_qubit_results.json and square_qubit_plot.png next to this script.
"""
import argparse, json, math, sys, time
from pathlib import Path
import numpy as np

HERE = Path(__file__).resolve().parent
PHASES_DEG = (0, 90, 180)
TIMES = np.linspace(0, math.pi / 2, 9)
SHOTS = 4000


# ----------------------------------------------------------------------------- maths
def prediction(phi_deg, t):
    return math.cos(math.radians(phi_deg) / 2) ** 2 * math.sin(2 * t) ** 2


def numpy_run(phi_deg, t):
    """Exact two-qubit statevector of the circuit, without Qiskit. Returns P(empty corners)."""
    phi = math.radians(phi_deg)
    psi = np.zeros(4, complex)                     # basis order |q1 q0>: 00, 01, 10, 11
    psi[0] = 1 / math.sqrt(2); psi[3] = np.exp(1j * phi) / math.sqrt(2)
    rx = lambda th: np.array([[math.cos(th / 2), -1j * math.sin(th / 2)], [-1j * math.sin(th / 2), math.cos(th / 2)]])
    U = np.kron(rx(-2 * t), rx(-2 * t))
    out = U @ psi
    p = np.abs(out) ** 2
    return p[1] + p[2]                             # |01> and |10> are beads 28 and 82


def selfcheck():
    worst = 0
    for phi in PHASES_DEG:
        for t in TIMES:
            worst = max(worst, abs(numpy_run(phi, t) - prediction(phi, t)))
    # perfect state transfer: |00> -> |11> at t = pi/2
    rx = lambda th: np.array([[math.cos(th / 2), -1j * math.sin(th / 2)], [-1j * math.sin(th / 2), math.cos(th / 2)]])
    psi = np.zeros(4, complex); psi[0] = 1
    out = np.kron(rx(-math.pi), rx(-math.pi)) @ psi
    pst = abs(out[3]) ** 2
    print(f"self-check: max |numpy - prediction| = {worst:.2e}; PST probability at t=pi/2 = {pst:.6f}")
    return worst < 1e-12 and abs(pst - 1) < 1e-12


# ----------------------------------------------------------------------------- circuits
def build_circuits():
    from qiskit import QuantumCircuit
    circuits, labels = [], []
    for phi in PHASES_DEG:
        for t in TIMES:
            qc = QuantumCircuit(2, 2)
            qc.h(0); qc.cx(0, 1)                    # (|00> + |11>)/sqrt2
            if phi:
                qc.rz(math.radians(phi), 1)         # relative phase phi on |11>
            qc.rx(-2 * t, 0); qc.rx(-2 * t, 1)      # exp(-iHt), H = -(X0 + X1)
            qc.measure([0, 1], [0, 1])
            circuits.append(qc); labels.append(("two_places", phi, float(t)))
    qc = QuantumCircuit(2, 2)                       # perfect state transfer, bead 1 -> bead 55
    qc.rx(-math.pi, 0); qc.rx(-math.pi, 1)
    qc.measure([0, 1], [0, 1])
    circuits.append(qc); labels.append(("pst", None, math.pi / 2))
    return circuits, labels


def summarise(counts, label):
    n = sum(counts.values())
    p = {k: v / n for k, v in counts.items()}
    kind, phi, t = label
    if kind == "pst":
        return {"kind": kind, "t": t, "P_bead55": p.get("11", 0.0), "counts": counts}
    return {"kind": kind, "phi_deg": phi, "t": t,
            "P_empty": p.get("01", 0.0) + p.get("10", 0.0), "prediction": prediction(phi, t), "counts": counts}


def run_ideal():
    from qiskit.primitives import StatevectorSampler
    circuits, labels = build_circuits()
    res = StatevectorSampler(default_shots=SHOTS).run(circuits).result()
    return [summarise(res[i].data.c.get_counts(), labels[i]) for i in range(len(circuits))], "statevector simulator"


def run_hardware(backend_name=None):
    from qiskit.transpiler import generate_preset_pass_manager
    from qiskit_ibm_runtime import QiskitRuntimeService, SamplerV2 as Sampler
    service = QiskitRuntimeService()
    backend = service.backend(backend_name) if backend_name else service.least_busy(operational=True, simulator=False, min_num_qubits=2)
    print("backend:", backend.name)
    circuits, labels = build_circuits()
    pm = generate_preset_pass_manager(backend=backend, optimization_level=3)
    isa = [pm.run(c) for c in circuits]
    sampler = Sampler(mode=backend)
    job = sampler.run(isa, shots=SHOTS)
    print("job id:", job.job_id(), "... waiting")
    t0 = time.time(); res = job.result(); print(f"done in {time.time() - t0:.0f} s")
    rows = [summarise(res[i].data.c.get_counts(), labels[i]) for i in range(len(circuits))]
    return rows, backend.name


# ----------------------------------------------------------------------------- report
def report(rows, source):
    out = {"source": source, "shots": SHOTS, "rows": rows}
    (HERE / "square_qubit_results.json").write_text(json.dumps(out, indent=1))
    pst = [r for r in rows if r["kind"] == "pst"][0]
    print(f"\nPerfect state transfer, bead 1 -> bead 55 at t = pi/2:  P = {pst['P_bead55']:.4f}   (prediction 1.0000)")
    print("\nTwo places at once, P(empty corners 28 + 82):")
    print("  phase    t      measured   predicted")
    for r in rows:
        if r["kind"] == "two_places":
            print(f"  {r['phi_deg']:>4}   {r['t']:.3f}    {r['P_empty']:.4f}     {r['prediction']:.4f}")
    try:
        import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
        S, T1, T2 = "#fcfcfb", "#0b0b0b", "#52514e"; cols = {0: "#2a78d6", 90: "#1baf7a", 180: "#eb6834"}
        fig, ax = plt.subplots(figsize=(8, 4.8), facecolor=S); ax.set_facecolor(S)
        tt = np.linspace(0, math.pi / 2, 200)
        for phi in PHASES_DEG:
            ax.plot(tt, [prediction(phi, t) for t in tt], ":", color=cols[phi], lw=1.3)
            pts = [r for r in rows if r["kind"] == "two_places" and r["phi_deg"] == phi]
            ax.plot([r["t"] for r in pts], [r["P_empty"] for r in pts], "o-", color=cols[phi], lw=2, ms=6, label=f"{phi}° between the two places")
        ax.set_title(f"Square qubit on {source}: probability on the empty corners", loc="left", fontsize=12, color=T1, weight="bold")
        ax.set_xlabel("evolution time t (dotted = cos²(φ/2)·sin²(2t))", color=T2); ax.set_ylabel("P(bead 28) + P(bead 82)", color=T2)
        ax.set_ylim(-0.02, 1.05); ax.legend(frameon=False, fontsize=9, labelcolor=T2); ax.tick_params(colors=T2)
        for s in ("top", "right"): ax.spines[s].set_visible(False)
        ax.text(0.02, 0.05, f"perfect state transfer check: P(bead 55) = {pst['P_bead55']:.3f} at t = π/2", transform=ax.transAxes, fontsize=9, color=T2)
        fig.savefig(HERE / "square_qubit_plot.png", dpi=130, facecolor=S, bbox_inches="tight")
        print("\nwrote square_qubit_results.json and square_qubit_plot.png")
    except Exception as e:
        print("plot skipped:", e)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--selfcheck", action="store_true"); ap.add_argument("--ideal", action="store_true")
    ap.add_argument("--hardware", action="store_true"); ap.add_argument("--backend", default=None)
    a = ap.parse_args()
    if a.selfcheck or not (a.ideal or a.hardware):
        sys.exit(0 if selfcheck() else 1)
    if a.ideal:
        rows, src = run_ideal(); report(rows, src)
    if a.hardware:
        rows, src = run_hardware(a.backend); report(rows, src)