#!/usr/bin/env python3
"""One engine, one job: the adapters that put a corpus circuit to each engine under test.

    python3 -I engines.py ENGINE JOB.json ANSWER.json

ENGINE is one of ADAPTERS below, named `<engine>.<method>`. Each runs in a process of its own,
so its peak memory is its own. A job carries the circuit's steps (corpus.py), its width and what
is asked:

  state        the whole state, written to job["state_out"] as complex128 pairs, qubit 0 the
               lowest bit of the index
  counts       job["shots"] samples of the classical bits at job["seed"]: {bits, bit 0 last: n}
  amplitudes   the amplitudes of job["bitstrings"] (qubit 0 last, as OpenQASM writes them)
  stabilizers  the stabilizer group, as signed strings of XYZ_ with qubit 0 first
  expectation  job["observable"] under job["noise"]
  qec_model    a circuit text's detector error model: [[p, [detectors], [observables]]]
  qec_sample   how often each detector and observable fired over job["shots"]

The third-party engines take each gate by their own name for it, where they have one, so a
gate an engine defines differently shows as a difference. Where an engine has no such gate, the
adapter says so in `notes` and gives it the gate's matrix from exact.py instead. An answer also
names what the engine was set to (precision, fusion, threads), the seconds its own call took
(`seconds`; for Aer the simulator's own time for the experiment, its whole `run` call beside it as
`call_seconds`), and the peak resident memory of the process that ran the engine, as the
operating system counts it for that process alone (`engine_peak_rss_bytes`: on Linux VmHWM,
which a child does not inherit from its parent the way ru_maxrss does). A job an engine cannot run
is answered {"ok": false, "unsupported": why}.

A job may carry `inject`, one of compare.py's faults for the engine under test: the adapter then
refuses, crashes, hangs or moves the angles before the engine reads them, as compare.py asks.
"""
import json
import os
import struct
import subprocess
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import exact  # noqa: E402


class Unsupported(Exception):
    pass


def threads(job):
    return int(job.get("threads") or 1)


def steps_of(job):
    return job["steps"]


def gates_only(job):
    for s in job["steps"]:
        if "g" not in s:
            raise Unsupported("this adapter runs programs with no measurement in the middle")
    return job["steps"]


def matrix_np(name, args):
    import numpy as np

    m = exact.one_qubit_matrix(name, args)
    s = float(exact.ONE)
    return np.array([[complex(m[0][0] / s, m[0][1] / s), complex(m[1][0] / s, m[1][1] / s)], [complex(m[2][0] / s, m[2][1] / s), complex(m[3][0] / s, m[3][1] / s)]])


def write_state(path, amps_le):
    """complex128 pairs, little-endian, qubit 0 the lowest index bit."""
    import numpy as np

    np.ascontiguousarray(amps_le, dtype=np.complex128).tofile(path)


def reverse_qubits(vec, n):
    """A state with qubit 0 as the highest index bit, reordered to qubit 0 lowest."""
    import numpy as np

    return np.ascontiguousarray(np.asarray(vec).reshape([2] * n).transpose(list(range(n - 1, -1, -1))).reshape(-1)) if n > 1 else np.asarray(vec)


# ------------------------------------------------------------------ the engine under test


def own_peak_rss():
    """This process's peak resident set, in bytes: VmHWM on Linux (its own address space only),
    ru_maxrss on macOS (bytes there)."""
    if sys.platform.startswith("linux"):
        for line in open("/proc/self/status"):
            if line.startswith("VmHWM:"):
                return int(line.split()[1]) * 1024
    import resource

    r = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    return r if sys.platform == "darwin" else r * 1024


def wai_driver(job_kind, payload, inject=None):
    exe = os.environ.get("WAI_CROSS_SIMULATOR")
    if not exe:
        raise Unsupported("WAI_CROSS_SIMULATOR names no driver")
    import tempfile

    if inject == "refuse":
        return {"ok": False, "refused": "an injected refusal"}
    if inject == "crash":
        raise RuntimeError("the driver exited 101: an injected crash")
    if inject == "hang" and job_kind == "qec_sample":
        time.sleep(10**6)
    if inject == "qec_words" and job_kind == "qec_model":
        return {"ok": False, "refused": "line 3: an injected refusal"}
    if inject == "qec_crash" and job_kind == "qec_model":
        raise RuntimeError("the driver exited 101: an injected crash")
    if inject == "exact_words" and job_kind == "state":
        return {"ok": False, "refused": "an injected refusal, in other words"}
    payload = dict(payload, kind=job_kind)
    with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f:
        json.dump(payload, f)
        path = f.name
    with tempfile.TemporaryFile() as so, tempfile.TemporaryFile() as se:
        try:
            p = subprocess.Popen([exe, path], stdout=so, stderr=se)
            # wait4 gives the driver's own peak on macOS; on Linux the driver reports VmHWM itself.
            _, status, ru = os.wait4(p.pid, 0)
        finally:
            os.unlink(path)
        so.seek(0)
        se.seek(0)
        stdout, stderr = so.read().decode(errors="replace"), se.read().decode(errors="replace")
    code = os.waitstatus_to_exitcode(status)
    if code != 0:
        raise RuntimeError("the driver exited %d: %s" % (code, stderr[-2000:]))
    v = json.loads(stdout)
    peak = v.get("peak_rss_bytes") or (ru.ru_maxrss if sys.platform == "darwin" else None)
    if not v["ok"]:
        return {"ok": False, "refused": v["refused"]}
    return {"ok": True, **v["answer"], "engine_peak_rss_bytes": peak}


def rounded_angles(steps):
    """Every angle moved to the nearest m/2^32 of a turn: what a reader that rounds would read."""
    from fractions import Fraction

    out = []
    for s in steps:
        if "a" in s:
            new = []
            for a in s["a"]:
                kind, v = exact.parse_angle(a)
                turns = v if kind == "t" else v / (2 * Fraction(exact._PI, 1 << (exact.PREC + exact.GUARD)))
                new.append("t:%d/32" % round(turns * (1 << 32)))
            s = dict(s, a=new)
        out.append(s)
    return out


def wai_state(job, method, controls=None):
    from corpus import qasm_of

    steps = gates_only(job)
    if job.get("inject") == "accept":
        steps = rounded_angles(steps)
    q = qasm_of({"n": job["n"], "bits": 0, "steps": steps})
    if job.get("perf"):
        # Timed: the state is made and not written.
        return wai_driver("state", {"qasm": q, "method": method, "fusion": job.get("fusion", 3), "ops": False, "controls": controls}, job.get("inject"))
    raw = job["state_out"] + ".raw"
    a = wai_driver("state", {"qasm": q, "method": method, "state_out": raw, "fusion": job.get("fusion", 3), "ops": method == "exact", "controls": controls}, job.get("inject"))
    if not a["ok"]:
        return a
    import numpy as np

    if method == "exact":
        ints = np.fromfile(raw, dtype="<i8").reshape(-1, 2)
        os.unlink(raw)
        # The integers themselves, beside the state, for the exact comparison.
        np.save(job["state_out"] + ".fixed.npy", ints)
        vec = (ints[:, 0] / float(1 << 30)) + 1j * (ints[:, 1] / float(1 << 30))
        a["settings"] = {"arithmetic": "fixed point, 2^-30", "fixed": job["state_out"] + ".fixed.npy"}
    else:
        vec = np.fromfile(raw, dtype="<f8").view(np.complex128)
        os.unlink(raw)
        a["settings"] = {"arithmetic": "double", "fusion": a.get("fusion")}
    write_state(job["state_out"], vec)
    return a


def wai_job(job, method):
    """Through the engine's whole path (plan, run, answers), as the job service runs a job."""
    from corpus import qasm_of

    q = qasm_of({"n": job["n"], "bits": job.get("bits", 0), "steps": job["steps"]})
    req = {"qasm": q, "method": method}
    kind = job["kind"]
    if kind == "counts":
        req["shots"] = job["shots"]
        req["seed"] = job.get("seed", 1)
    elif kind == "amplitudes":
        req["amplitudes"] = job["bitstrings"]
        req["shots"] = 0
    elif kind == "expectation":
        req["observables"] = [job["observable"]]
        req["noise"] = job.get("noise", {})
        req["shots"] = 0
    if job.get("options"):
        req["options"] = job["options"]
    a = wai_driver("job", {"request": req, "memory_gb": job.get("memory_gb", 16)}, job.get("inject"))
    if not a["ok"]:
        return a
    r = a["result"]
    out = {"ok": True, "seconds": a["seconds"], "planned": a["method"], "accuracy": r.get("accuracy"), "engine_peak_rss_bytes": a.get("engine_peak_rss_bytes")}
    if kind == "counts":
        out["counts"] = r.get("counts", {})
        if "outcome_dimension" in r:
            out["outcome_dimension"] = r["outcome_dimension"]
    elif kind == "amplitudes":
        out["amplitudes"] = {x["bits"]: [x["re"], x["im"]] for x in r["amplitudes"]}
        if "contraction_flops" in r:
            out["contraction_flops"] = r["contraction_flops"]
    elif kind == "expectation":
        out["value"] = r["expectations"][0]["value"]
    return out


def wai_stabilizers(job):
    from corpus import qasm_of

    q = qasm_of({"n": job["n"], "bits": 0, "steps": gates_only(job)})
    return wai_driver("stabilizers", {"qasm": q, "paulis": job["paulis"]}, job.get("inject"))


def wai_qec(job):
    if job["kind"] == "qec_model":
        a = wai_driver("qec_model", {"text": job["text"]}, job.get("inject"))
        if not a["ok"]:
            return a
        mech = []
        for line in a["model"].splitlines():
            line = line.strip()
            if not line.startswith("error("):
                continue
            p, _, rest = line[6:].partition(")")
            ds, ls = [], []
            for tok in rest.split():
                if tok.startswith("D"):
                    ds.append(int(tok[1:]))
                elif tok.startswith("L"):
                    ls.append(int(tok[1:]))
            mech.append([float(p), sorted(ds), sorted(ls)])
        return {"ok": True, "seconds": a["seconds"], "detectors": a["detectors"], "observables": a["observables"], "mechanisms": mech, "engine_peak_rss_bytes": a.get("engine_peak_rss_bytes")}
    return wai_driver("qec_sample", {"text": job["text"], "shots": job["shots"], "seed": job.get("seed", 1)}, job.get("inject"))


# ------------------------------------------------------------------ Qiskit circuits (Aer, DDSIM)


def qiskit_circuit(job, measure=True):
    from qiskit import QuantumCircuit
    from qiskit.circuit.library import GlobalPhaseGate, PhaseGate, RYGate, RZGate, XGate, ZGate, UnitaryGate

    n, nb = job["n"], job.get("bits", 0)
    qc = QuantumCircuit(n, nb) if nb else QuantumCircuit(n)
    notes = set()

    def gate(s, target=qc):
        name, q, args = s["g"], s["q"], [exact.float_radians(a) for a in s.get("a", [])]
        simple = {"id": "id", "x": "x", "y": "y", "z": "z", "h": "h", "s": "s", "sdg": "sdg", "t": "t", "tdg": "tdg", "sx": "sx", "sxdg": "sxdg", "p": "p", "rz": "rz", "rx": "rx", "ry": "ry", "u": "u", "cx": "cx", "cy": "cy", "cz": "cz", "ch": "ch", "cp": "cp", "crx": "crx", "cry": "cry", "crz": "crz", "cu": "cu", "csx": "csx", "swap": "swap", "cswap": "cswap", "ccx": "ccx", "ccz": "ccz", "rzz": "rzz", "rxx": "rxx", "ryy": "ryy"}
        if name in simple:
            getattr(target, simple[name])(*args, *q)
        elif name == "negcx":
            target.append(XGate().control(1, ctrl_state=0), q)
        elif name == "gphase":
            if q:
                target.append(GlobalPhaseGate(args[0]).control(len(q)), q)
            else:
                target.global_phase += args[0]
        else:
            form = exact.controlled_form(name)
            if not form:
                raise Unsupported("no gate %s" % name)
            base, k = form
            g = {"x": XGate(), "z": ZGate(), "p": PhaseGate(*args) if args else None, "rz": RZGate(*args) if args else None, "ry": RYGate(*args) if args else None}.get(base)
            if g is None:
                g = UnitaryGate(matrix_np(base, s.get("a", [])))
                notes.add("%s as a matrix" % name)
            target.append(g.control(k), q)

    for s in job["steps"]:
        if "g" in s:
            gate(s)
        elif "measure" in s:
            if measure:
                qc.measure(s["measure"], s["bit"])
        elif "reset" in s:
            qc.reset(s["reset"])
        elif "if" in s:
            c, v = s["if"]
            with qc.if_test((qc.clbits[c], v)):
                for x in s["then"]:
                    gate(x)
    return qc, sorted(notes)


def aer_run(job, method):
    from qiskit import transpile
    from qiskit_aer import AerSimulator

    kind = job["kind"]
    opts = {"method": method, "max_parallel_threads": threads(job), "seed_simulator": job.get("seed", 1)}
    if method == "matrix_product_state" and not job.get("perf"):
        # Its log names every two-site update, the bonds after it and each truncation's weight
        # (not kept in a timed run, whose time it would add to).
        opts["mps_log_data"] = True
    if job.get("aer_options"):
        opts.update(job["aer_options"])
    sim = AerSimulator(**opts)
    if kind == "qasm_state":
        from qiskit import qasm3

        qc = qasm3.loads(job["qasm"])
        notes = ["read from the published OpenQASM 3 text"]
        qc.save_statevector()
    else:
        qc, notes = qiskit_circuit(job)
        if kind == "state" and job.get("perf"):
            # Timed: the whole state is made; one amplitude is handed back, not the state.
            gates_only(job)
            qc.save_amplitudes([0])
        elif kind == "state":
            gates_only(job)
            qc.save_statevector()
        elif kind == "stabilizers" and job.get("perf"):
            # Timed: the gates on the tableau, and one qubit measured. Reading the whole group
            # out through the save instruction's result is not the simulation, and is left out.
            from qiskit import ClassicalRegister

            gates_only(job)
            qc.add_register(ClassicalRegister(1))
            qc.measure(0, 0)
        elif kind == "stabilizers":
            gates_only(job)
            qc.save_stabilizer()
        elif kind == "expectation":
            from qiskit.quantum_info import Pauli

            label = ["I"] * job["n"]
            for f in job["observable"].split():
                label[int(f[1:])] = f[0]
            qc.save_expectation_value(Pauli("".join(reversed(label))), list(range(job["n"])))
        elif kind == "amplitudes":
            qc.save_amplitudes([int(b, 2) for b in job["bitstrings"]])
    if kind == "expectation" and job.get("noise"):
        from qiskit_aer.noise import NoiseModel, depolarizing_error

        nm = NoiseModel()
        for k, p in job["noise"].items():
            if k == "depolarizing_1q":
                nm.add_all_qubit_quantum_error(depolarizing_error(4 * p / 3, 1), ["x", "h", "s", "sdg", "t", "tdg", "sx", "rz", "rx", "ry", "p", "u"])
            elif k == "depolarizing_2q":
                nm.add_all_qubit_quantum_error(depolarizing_error(16 * p / 15, 2), ["cx", "cz"])
            else:
                raise Unsupported("noise %s" % k)
        sim = AerSimulator(noise_model=nm, **opts)
    tqc = transpile(qc, sim, optimization_level=0)
    t0 = time.perf_counter()
    res = sim.run(tqc, shots=job.get("shots") or (1 if kind != "counts" else 1024)).result()
    secs = time.perf_counter() - t0
    if not res.success:
        raise RuntimeError(res.status)
    data = res.data(0)
    settings = {"method": method, "precision": sim.options.precision, "threads": threads(job), "fusion_threshold": 14, "fusion_max_qubit": 5, "gates": len(tqc.data)}
    # `seconds`: the simulator's own time for the experiment, as its result states it; the
    # whole call, with the conversion of the circuit to its own form, is `call_seconds`.
    out = {"ok": True, "seconds": res.results[0].time_taken, "call_seconds": secs, "settings": settings, "notes": notes}
    if method == "matrix_product_state" and not job.get("perf"):
        import re

        log = str(res.results[0].metadata.get("MPS_log_data", ""))
        bonds = [int(x) for b in re.findall(r"BD=\[([0-9 ]*)\]", log) for x in b.split()]
        settings["mps_log"] = {"updates": len(re.findall(r"BD=\[", log)), "max_bond": max(bonds, default=1), "discarded": [float(x) for x in re.findall(r"discarded_value=([0-9.eE+-]+)", log)]}
    if kind in ("state", "qasm_state") and not job.get("perf"):
        import numpy as np

        write_state(job["state_out"], np.asarray(data["statevector"]))
    elif kind == "stabilizers" and job.get("perf"):
        pass
    elif kind == "stabilizers":
        cl = data["stabilizer"]
        labels = cl.to_labels(mode="S") if hasattr(cl, "to_labels") else cl.clifford.to_labels(mode="S")
        rows = []
        for lab in labels:
            sign = "-" if lab.startswith("-") else "+"
            body = lab.lstrip("+-")
            rows.append(sign + "".join(reversed(body)).replace("I", "_"))
        out["stabilizers"] = rows
    elif kind == "counts":
        out["counts"] = {k.replace(" ", ""): v for k, v in res.get_counts(0).items()}
    elif kind == "expectation":
        out["value"] = float(data["expectation_value"])
    elif kind == "amplitudes":
        out["amplitudes"] = {b: [complex(a).real, complex(a).imag] for b, a in zip(job["bitstrings"], data["amplitudes"])}
    return out


def ddsim_run(job, backend_name):
    from mqt.ddsim import DDSIMProvider

    qc, notes = qiskit_circuit(job)
    kind = job["kind"]
    if kind == "state":
        gates_only(job)
    b = DDSIMProvider().get_backend(backend_name)
    t0 = time.perf_counter()
    if kind == "counts":
        res = b.run(qc, shots=job["shots"], seed=job.get("seed", 1)).result()
    else:
        res = b.run(qc).result()
    secs = time.perf_counter() - t0
    out = {"ok": True, "seconds": secs, "settings": {"arithmetic": "decision diagram, double", "tolerance": 2.0**-52 * 1024, "parameter_tolerance": 1e-13}, "notes": notes}
    if kind == "state" and job.get("perf"):
        pass  # timed: the run made the state vector, and it is not written
    elif kind == "state":
        import numpy as np

        write_state(job["state_out"], np.asarray(res.get_statevector(qc)))
    else:
        out["counts"] = {k.replace(" ", ""): v for k, v in res.get_counts(qc).items()}
    return out


# ------------------------------------------------------------------ quimb


QUIMB_NAMES = {"id": "IDEN", "x": "X", "y": "Y", "z": "Z", "h": "H", "s": "S", "sdg": "SDG", "t": "T", "tdg": "TDG", "sx": "SX", "sxdg": "SXDG", "rz": "RZ", "rx": "RX", "ry": "RY", "u": "U3", "p": "U1", "cx": "CX", "cy": "CY", "cz": "CZ", "cp": "CU1", "crx": "CRX", "cry": "CRY", "crz": "CRZ", "swap": "SWAP", "cswap": "CSWAP", "ccx": "CCX", "ccz": "CCZ", "rzz": "RZZ", "rxx": "RXX", "ryy": "RYY"}


def quimb_circuit(job, cls, **kw):
    import quimb.tensor as qtn

    n = job["n"]
    c = getattr(qtn, cls)(n, **kw)
    notes = set()
    for s in gates_only(job):
        name, q, args = s["g"], s["q"], [exact.float_radians(a) for a in s.get("a", [])]
        if name in QUIMB_NAMES:
            c.apply_gate(QUIMB_NAMES[name], *args, *q)
        elif name in ("ch", "csx"):
            c.apply_gate({"ch": "H", "csx": "SX"}[name], q[1], controls=[q[0]])
        elif name == "cu":
            # stdgates.inc: p(γ) on the control, then the controlled U.
            c.apply_gate("U1", args[3], q[0])
            c.apply_gate("U3", args[0], args[1], args[2], q[1], controls=[q[0]])
        elif name == "negcx":
            c.apply_gate("X", q[0])
            c.apply_gate("CX", q[0], q[1])
            c.apply_gate("X", q[0])
            notes.add("negcx as x, cx, x")
        else:
            form = exact.controlled_form(name)
            if not form:
                raise Unsupported("no gate %s" % name)
            base, k = form
            qn = {"x": "X", "z": "Z", "p": "U1", "rz": "RZ", "ry": "RY", "h": "H", "sx": "SX", "y": "Y"}[base]
            c.apply_gate(qn, *args, q[k], controls=list(q[:k]))
    return c, sorted(notes)


def quimb_run(job, how):
    import numpy as np

    kind = job["kind"]
    n = job["n"]
    if how == "dense":
        # The dense form applies each gate as it is added: the time is the circuit's.
        t0 = time.perf_counter()
        c, notes = quimb_circuit(job, "CircuitDense")
        vec = np.asarray(c.psi.to_dense()).reshape(-1)
        secs = time.perf_counter() - t0
        settings = {"arithmetic": "double", "form": "dense, gate by gate"}
    elif how == "mps":
        # The matrix product state is updated as each gate is added: the time starts here.
        t_build = time.perf_counter()
        # An absolute cutoff: each update drops singular values whose squares sum to under 1e-28,
        # so what it can drop is known without its log, and a double's rounding (squares near
        # 1e-32) is still dropped.
        c, notes = quimb_circuit(job, "CircuitMPS", gate_opts={"cutoff": 1e-28, "cutoff_mode": "sum2", "max_bond": job.get("max_bond")})
        settings = {"arithmetic": "double", "form": "matrix product state", "cutoff": 1e-28, "cutoff_mode": "sum2", "max_bond": job.get("max_bond")}
        secs = 0.0
        vec = None
    else:
        t_build = time.perf_counter()
        c, notes = quimb_circuit(job, "Circuit")
        settings = {"arithmetic": "double", "form": "tensor network contracted"}
        secs = 0.0
        vec = None
    out = {"ok": True, "settings": settings, "notes": notes, "gates": len(c.gates)}
    if kind == "state":
        if vec is None:
            t0 = time.perf_counter()
            vec = np.asarray(c.to_dense()).reshape(-1)
            secs = time.perf_counter() - t0
            if how == "tn":
                settings["contraction_width"] = float(c.psi.contraction_width())
                settings["contraction_steps"] = len(c.psi.tensors) - 1
        if not job.get("perf"):
            write_state(job["state_out"], reverse_qubits(vec, n))
    elif kind == "amplitudes":
        amps = {}
        t0 = time.perf_counter()
        widths = []
        for b in job["bitstrings"]:
            # quimb's bit string has qubit 0 first; OpenQASM's has it last.
            qb = b[::-1]
            if how == "tn":
                tn = c.amplitude_tn(qb)
                widths.append(tn.contraction_width())
                a = complex(tn.contract(all, optimize="auto-hq"))
            else:
                a = complex(c.amplitude(qb)) if how != "mps" else complex(c.psi.amplitude(qb))
            amps[b] = [a.real, a.imag]
        secs = time.perf_counter() - t_build
        if widths:
            settings["contraction_width"] = max(widths)
            settings["contraction_steps"] = len(c.psi.tensors)
        if how == "mps":
            settings["max_bond_used"] = int(c.psi.max_bond())
            settings["fidelity_estimate"] = float(c.fidelity_estimate())
        out["amplitudes"] = amps
    else:
        raise Unsupported("quimb answers states and amplitudes here")
    out["seconds"] = secs
    return out


# ------------------------------------------------------------------ qsim, through Cirq


def cirq_circuit(job):
    import cirq
    import numpy as np

    n = job["n"]
    qs = cirq.LineQubit.range(n)
    ops = []
    notes = set()
    for s in gates_only(job):
        name, q, a = s["g"], [qs[i] for i in s["q"]], [exact.float_radians(x) for x in s.get("a", [])]
        pi = np.pi
        simple = {"id": cirq.I, "x": cirq.X, "y": cirq.Y, "z": cirq.Z, "h": cirq.H, "s": cirq.S, "sdg": cirq.S**-1, "t": cirq.T, "tdg": cirq.T**-1, "sx": cirq.X**0.5, "sxdg": cirq.X**-0.5, "cx": cirq.CNOT, "cz": cirq.CZ, "swap": cirq.SWAP, "ccx": cirq.CCX, "ccz": cirq.CCZ, "cswap": cirq.CSWAP}
        if name in simple:
            ops.append(simple[name].on(*q))
        elif name == "rz":
            ops.append(cirq.rz(a[0]).on(*q))
        elif name == "rx":
            ops.append(cirq.rx(a[0]).on(*q))
        elif name == "ry":
            ops.append(cirq.ry(a[0]).on(*q))
        elif name == "p":
            ops.append(cirq.ZPowGate(exponent=a[0] / pi).on(*q))
        elif name == "cp":
            ops.append(cirq.CZPowGate(exponent=a[0] / pi).on(*q))
        elif name == "rzz":
            ops.append(cirq.ZZPowGate(exponent=a[0] / pi, global_shift=-0.5).on(*q))
        elif name == "rxx":
            ops.append(cirq.XXPowGate(exponent=a[0] / pi, global_shift=-0.5).on(*q))
        elif name == "ryy":
            ops.append(cirq.YYPowGate(exponent=a[0] / pi, global_shift=-0.5).on(*q))
        elif name in ("cy", "ch", "csx", "crx", "cry", "crz"):
            base = {"cy": cirq.Y, "ch": cirq.H, "csx": cirq.X**0.5, "crx": cirq.rx(a[0]) if a else None, "cry": cirq.ry(a[0]) if a else None, "crz": cirq.rz(a[0]) if a else None}[name]
            ops.append(base.controlled(1).on(*q))
        elif name == "negcx":
            ops.append(cirq.X.controlled(1, control_values=[0]).on(*q))
        elif name in ("u", "cu"):
            m = matrix_np("u", s["a"][:3])
            if name == "u":
                ops.append(cirq.MatrixGate(m).on(*q))
            else:
                ops.append(cirq.ZPowGate(exponent=a[3] / pi).on(q[0]))
                ops.append(cirq.MatrixGate(m).controlled(1).on(*q))
            notes.add("U as a matrix")
        elif name == "gphase":
            raise Unsupported("a global phase gate")
        else:
            form = exact.controlled_form(name)
            if not form:
                raise Unsupported("no gate %s" % name)
            base, k = form
            g = {"x": cirq.X, "z": cirq.Z, "p": cirq.ZPowGate(exponent=a[0] / pi) if a else None, "rz": cirq.rz(a[0]) if a else None, "ry": cirq.ry(a[0]) if a else None}[base]
            ops.append(g.controlled(k).on(*q))
    return cirq.Circuit(ops), qs, sorted(notes)


def qsim_run(job):
    import numpy as np
    import qsimcirq

    kind = job["kind"]
    if kind == "state":
        c, qs, notes = cirq_circuit(job)
        sim = qsimcirq.QSimSimulator(qsimcirq.QSimOptions(cpu_threads=threads(job), max_fused_gate_size=2))
        if job.get("perf"):
            # Timed: the whole state is made; one amplitude is handed back, not the state.
            t0 = time.perf_counter()
            sim.compute_amplitudes(c, bitstrings=[0], qubit_order=qs)
            return {"ok": True, "seconds": time.perf_counter() - t0, "settings": {"arithmetic": "single", "fusion": 2, "threads": threads(job)}}
        t0 = time.perf_counter()
        r = sim.simulate(c, qubit_order=qs)
        secs = time.perf_counter() - t0
        vec = np.asarray(r.final_state_vector)
        write_state(job["state_out"], reverse_qubits(vec.astype(np.complex128), job["n"]))
        return {"ok": True, "seconds": secs, "settings": {"arithmetic": "single", "fusion": 2, "threads": threads(job), "dtype": str(vec.dtype)}, "notes": notes}
    if kind == "counts":
        import cirq

        steps = job["steps"]
        k = len(steps)
        while k and "measure" in steps[k - 1]:
            k -= 1
        if any("g" not in s for s in steps[:k]) or len({m["measure"] for m in steps[k:]}) != len(steps) - k:
            raise Unsupported("measurement in the middle (or one qubit measured twice)")
        sub = dict(job, steps=steps[:k])
        c, qs, notes = cirq_circuit(sub)
        meas = steps[k:]
        c.append(cirq.measure(*[qs[m["measure"]] for m in meas], key="m"))
        sim = qsimcirq.QSimSimulator(qsimcirq.QSimOptions(cpu_threads=threads(job)), seed=job.get("seed", 1))
        t0 = time.perf_counter()
        r = sim.run(c, repetitions=job["shots"])
        secs = time.perf_counter() - t0
        rows = r.measurements["m"]
        nb = job["bits"]
        counts = {}
        for row in rows:
            bits = [0] * nb
            for m, v in zip(meas, row):
                bits[m["bit"]] = int(v)
            key = "".join(str(bits[i]) for i in range(nb - 1, -1, -1))
            counts[key] = counts.get(key, 0) + 1
        return {"ok": True, "seconds": secs, "counts": counts, "settings": {"arithmetic": "single"}, "notes": notes}
    raise Unsupported("qsim answers states and counts here")


# ------------------------------------------------------------------ Stim


STIM_NAMES = {"id": "I", "x": "X", "y": "Y", "z": "Z", "h": "H", "s": "S", "sdg": "S_DAG", "sx": "SQRT_X", "sxdg": "SQRT_X_DAG", "cx": "CX", "cy": "CY", "cz": "CZ", "swap": "SWAP"}


def stim_text(job):
    lines = []
    last_record = {}
    records = 0
    for s in job["steps"]:
        if "g" in s:
            if s["g"] not in STIM_NAMES:
                raise Unsupported("%s is not a Clifford gate" % s["g"])
            lines.append("%s %s" % (STIM_NAMES[s["g"]], " ".join(map(str, s["q"]))))
        elif "measure" in s:
            lines.append("M %d" % s["measure"])
            records += 1
            last_record[s["bit"]] = records
        elif "reset" in s:
            lines.append("R %d" % s["reset"])
        elif "if" in s:
            c, v = s["if"]
            if v != 1 or c not in last_record:
                raise Unsupported("feedback on a bit's 0, or on a bit not yet measured")
            for x in s["then"]:
                op = {"x": "CX", "z": "CZ", "y": "CY"}.get(x["g"])
                if not op:
                    raise Unsupported("feedback by %s" % x["g"])
                lines.append("%s rec[-%d] %d" % (op, records - last_record[c] + 1, x["q"][0]))
    return "\n".join(lines) + "\n", last_record, records


def stim_run(job):
    import stim

    kind = job["kind"]
    if kind == "stabilizers":
        text, _, _ = stim_text(dict(job, steps=gates_only(job)))
        c = stim.Circuit(text)
        t0 = time.perf_counter()
        sim = stim.TableauSimulator()
        sim.set_num_qubits(job["n"])
        sim.do_circuit(c)
        secs = time.perf_counter() - t0
        if job.get("perf"):
            return {"ok": True, "seconds": secs}
        t1 = time.perf_counter()
        stabs = sim.canonical_stabilizers()
        return {"ok": True, "seconds": secs, "canonical_seconds": time.perf_counter() - t1, "stabilizers": [str(p) for p in stabs]}
    if kind == "counts":
        text, last_record, records = stim_text(job)
        c = stim.Circuit(text)
        t0 = time.perf_counter()
        samples = c.compile_sampler(seed=job.get("seed", 1)).sample(job["shots"])
        secs = time.perf_counter() - t0
        nb = job["bits"]
        counts = {}
        import numpy as np

        keys = np.zeros(job["shots"], dtype=object)
        cols = []
        for b in range(nb - 1, -1, -1):
            cols.append(samples[:, last_record[b] - 1].astype(np.uint8) if b in last_record else np.zeros(job["shots"], dtype=np.uint8))
        mat = np.stack(cols, axis=1) if cols else np.zeros((job["shots"], 0), dtype=np.uint8)
        packed, inverse = np.unique(mat, axis=0, return_inverse=True)
        tallies = np.bincount(inverse.reshape(-1), minlength=len(packed))
        for row, k in zip(packed, tallies):
            counts["".join(str(int(x)) for x in row)] = int(k)
        _ = keys
        return {"ok": True, "seconds": secs, "counts": counts}
    if kind == "qec_model":
        c = stim.Circuit(job["text"])
        t0 = time.perf_counter()
        dem = c.detector_error_model(decompose_errors=False, flatten_loops=True)
        secs = time.perf_counter() - t0
        mech = []
        for ins in dem.flattened():
            if ins.type != "error":
                continue
            ds, ls = [], []
            for tg in ins.targets_copy():
                if tg.is_relative_detector_id():
                    ds.append(tg.val)
                elif tg.is_logical_observable_id():
                    ls.append(tg.val)
            mech.append([ins.args_copy()[0], sorted(ds), sorted(ls)])
        return {"ok": True, "seconds": secs, "detectors": dem.num_detectors, "observables": dem.num_observables, "mechanisms": mech}
    if kind == "qec_sample":
        c = stim.Circuit(job["text"])
        if job.get("perf"):
            # Timed: the samples bit-packed, as the engine under test holds its own.
            t0 = time.perf_counter()
            c.compile_detector_sampler(seed=job.get("seed", 1)).sample(job["shots"], separate_observables=True, bit_packed=True)
            return {"ok": True, "seconds": time.perf_counter() - t0, "shots": job["shots"]}
        t0 = time.perf_counter()
        det, obs = c.compile_detector_sampler(seed=job.get("seed", 1)).sample(job["shots"], separate_observables=True)
        secs = time.perf_counter() - t0
        return {"ok": True, "seconds": secs, "shots": job["shots"], "detectors": [int(x) for x in det.sum(axis=0)], "observables": [int(x) for x in obs.sum(axis=0)]}
    raise Unsupported(kind)


def stim_generated(job):
    """A circuit from the pinned generator, as text."""
    import stim

    c = stim.Circuit.generated(job["code"], distance=job["distance"], rounds=job["rounds"], after_clifford_depolarization=job["p"], before_round_data_depolarization=job["p"], before_measure_flip_probability=job["p"], after_reset_flip_probability=job["p"])
    return {"ok": True, "text": str(c) + "\n"}


# ------------------------------------------------------------------ the table


ADAPTERS = {
    "wai.exact": lambda j: wai_state(j, "exact") if j["kind"] == "state" else wai_job(j, "exact"),
    # The bit-exact method asked to hold a gate under many controls whole ("controls": "whole"):
    # other bits than its default, which runs standard gates.
    "wai.exact_whole": lambda j: wai_state(j, "exact", "whole") if j["kind"] == "state" else wai_job(dict(j, options={"controls": "whole"}), "exact"),
    "wai.statevector": lambda j: wai_state(j, "statevector") if j["kind"] == "state" else wai_job(j, "statevector"),
    "wai.stabilizer": lambda j: wai_stabilizers(j) if j["kind"] == "stabilizers" else wai_job(j, "stabilizer"),
    "wai.auto": lambda j: wai_job(j, "auto"),
    "wai.mps": lambda j: wai_job(dict(j, options={"max_bond": j.get("max_bond") or 64}), "mps"),
    "wai.tensor_network": lambda j: wai_job(j, "tensor_network"),
    "wai.density_matrix": lambda j: wai_job(j, "density_matrix"),
    "wai.qec": wai_qec,
    "aer.statevector": lambda j: aer_run(j, "statevector"),
    "aer.qasm_import": lambda j: aer_run(dict(j, kind="qasm_state"), "statevector"),
    "aer.stabilizer": lambda j: aer_run(j, "stabilizer"),
    "aer.mps": lambda j: aer_run(dict(j, aer_options={"matrix_product_state_truncation_threshold": 1e-16}), "matrix_product_state"),
    "aer.density_matrix": lambda j: aer_run(j, "density_matrix"),
    "stim.tableau": stim_run,
    "stim.generate": stim_generated,
    "quimb.dense": lambda j: quimb_run(j, "dense"),
    "quimb.tensor_network": lambda j: quimb_run(j, "tn"),
    "quimb.mps": lambda j: quimb_run(j, "mps"),
    "ddsim.statevector": lambda j: ddsim_run(j, "statevector_simulator" if j["kind"] == "state" else "qasm_simulator"),
    "qsim.statevector": qsim_run,
}


def versions():
    """Each engine's version, as it reports it."""
    out = {}
    for name, mod in (("stim", "stim"), ("qiskit", "qiskit"), ("qiskit-aer", "qiskit_aer"), ("quimb", "quimb"), ("mqt.ddsim", "mqt.ddsim"), ("mqt.core", "mqt.core"), ("qsimcirq", "qsimcirq"), ("cirq-core", "cirq"), ("numpy", "numpy"), ("scipy", "scipy")):
        try:
            m = __import__(mod, fromlist=["__version__"])
            out[name] = m.__version__
        except Exception as e:  # noqa: BLE001 - reported, not hidden
            out[name] = "unavailable: %s" % type(e).__name__
    out["python"] = sys.version.split()[0]
    return out


def main():
    if sys.argv[1:2] == ["--versions"]:
        print(json.dumps(versions(), sort_keys=True))
        return
    engine, job_path, answer_path = sys.argv[1:4]
    job = json.load(open(job_path))
    try:
        answer = ADAPTERS[engine](job)
    except Unsupported as e:
        answer = {"ok": False, "unsupported": str(e)}
    except Exception as e:  # noqa: BLE001 - an engine's failure is a result, kept with its message
        answer = {"ok": False, "error": "%s: %s" % (type(e).__name__, str(e)[:2000])}
    # The adapter's own peak; for another engine, which runs in this process, the engine's too.
    answer["adapter_peak_rss_bytes"] = own_peak_rss()
    if not engine.startswith("wai."):
        answer["engine_peak_rss_bytes"] = answer["adapter_peak_rss_bytes"]
    with open(answer_path, "w") as f:
        json.dump(answer, f)


if __name__ == "__main__":
    main()