#!/usr/bin/env python3
"""gridsynth_oracle.py: the reference T counts that the engine's z-rotation
synthesiser (wai/wai-quantum/src/engine/synth/rz.rs) is held to.

For each angle and precision it runs pygridsynth 1.0.1, the reference
implementation of the Ross–Selinger algorithm that arXiv:2506.03094 cites, and
prints one CSV row: the angle and the precision as the doubles both sides use
(Python's repr, which Rust parses back to the same bits), pygridsynth's
denominator exponent and its T count. The output is
wai/wai-quantum/tests/data/rz_gridsynth_oracle.csv. CI does not run this
script and nothing depends on pygridsynth: the CSV is the oracle's answer,
committed, and `engine::synth::rz`'s tests read it.

    python3 -m venv /tmp/pgs && /tmp/pgs/bin/pip install "pygridsynth==1.0.1" mpmath
    /tmp/pgs/bin/python gridsynth_oracle.py > rz_gridsynth_oracle.csv

The angles:
- the four rotation angles of the 10 × 10 Ising program of arXiv:2506.03094
  (Appendix A.10: a fourth-order product formula, Δ = 0.25, J = g = 1), at
  its ε = 5·10⁻⁹ and at every precision below;
- 128 angles spread over (−π, π) by the golden ratio, θᵢ = −π + 2π·frac(i·φ),
  none within 10⁻⁶ of a multiple of π/4 (those the engine writes exactly);
at ε = 10⁻¹, 10⁻², 10⁻³, 10⁻⁴, 10⁻⁵, 10⁻⁷, 10⁻⁹, 10⁻¹¹ and 10⁻¹³.
"""
import io
import math
import multiprocessing
import contextlib
import sys

import mpmath
from pygridsynth.gridsynth import gridsynth_gates
from pygridsynth.gridsynth import gridsynth

PRECISIONS = [1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-7, 1e-9, 1e-11, 1e-13]


def ising_angles():
    gamma = 1.0 / (4.0 - 4.0 ** (1.0 / 3.0))
    delta = 0.25
    mid = (1.0 - 4.0 * gamma) * delta
    return [("ising zz", gamma * delta), ("ising x", 2.0 * gamma * delta), ("ising zz middle", mid), ("ising x middle", 2.0 * mid)]


def swept_angles(n=128):
    phi = (1.0 + math.sqrt(5.0)) / 2.0
    out = []
    i = 1
    while len(out) < n:
        theta = -math.pi + 2.0 * math.pi * math.fmod(i * phi, 1.0)
        q = theta / (math.pi / 4.0)
        if abs(q - round(q)) > 1e-6:
            out.append((f"sweep {len(out) + 1}", theta))
        i += 1
    return out


def run(job):
    name, theta, eps = job
    mpmath.mp.dps = 128
    with contextlib.redirect_stdout(io.StringIO()):
        u = gridsynth(theta=mpmath.mpf(theta), epsilon=mpmath.mpf(eps))
        gates = gridsynth_gates(theta=mpmath.mpf(theta), epsilon=mpmath.mpf(eps))
    return (name, repr(theta), repr(eps), u.k, gates.count("T"))


def main():
    jobs = [(name, theta, 5e-9) for name, theta in ising_angles()]
    for eps in PRECISIONS:
        jobs += [(name, theta, eps) for name, theta in ising_angles() + swept_angles()]
    with multiprocessing.Pool(int(sys.argv[1]) if len(sys.argv) > 1 else 8) as pool:
        rows = pool.map(run, jobs, chunksize=1)
    print("name,theta,epsilon,k,t_count")
    for row in rows:
        print(",".join(str(x) for x in row))


if __name__ == "__main__":
    main()