"""Emit the TikZ for the deviation-rate figure from the recorded run.

The figure plots, against the noise level, the measured probability that a step
of the levelized host deviates from the exact orbit and the union bound of the
random-perturbation proposition. Both axes are read from
paper/runs/paper_noise_curve.json.

    python src/make_figures.py > /tmp/curve.tex
"""
import json
import math
import os
import sys

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def load(name):
    return json.load(open(os.path.join(REPO, "paper", "runs", name)))


def group(n):
    return f"{n:,}".replace(",", "{,}")


def sci(x, digits=2):
    """LaTeX scientific notation."""
    if x == 0:
        return "$0$"
    e = math.floor(math.log10(abs(x)))
    m = x / 10 ** e
    if e == 0:
        return f"${m:.{digits}f}$"
    return rf"${m:.{digits}f}\times 10^{{{e}}}$"


def omega_table():
    d = load("paper_noise_omega.json")
    rows = []
    for r in d["rows"]:
        if r["complete"]:
            outcome = "emitted the whole instance"
        elif r["first_wrong_byte"] is not None:
            outcome = (f"a wrong byte after {group(r['first_wrong_byte'])} "
                       f"correct")
        elif r["halted"]:
            outcome = "the halt bit was set"
        else:
            outcome = "the budget was exhausted"
        rows.append(f"${r['sigma']:.2f}$ & {group(r['steps'])} & "
                    f"{group(r['bytes'])} & {outcome}\\\\")
    body = "\n".join(rows)
    return rf"""\begin{{table}}[ht]
\centering\small
\begin{{tabular}}{{@{{}}lrrl@{{}}}}
\toprule
$s$ & Steps executed & Bytes emitted & Outcome\\
\midrule
{body}
\bottomrule
\end{{tabular}}
\caption{{The self-reproducing instance $\Omega^\ast$ on $\Lev(\Nhost)$ with
additive Gaussian read noise of standard deviation $s$ applied independently to
every pre-activation of every layer at every step, thresholded at
$-\tfrac12$. A run ends when it emits a byte that differs from
$\ser(\Omega^\ast)$, when the halt bit is set, or when the whole instance has
been emitted. The target is {group(d['target_bytes'])} bytes.}}
\label{{tab:noise}}
\end{{table}}"""


def leak_table():
    d = load("paper_noise_leak.json")
    rows = []
    for r in d["rows"]:
        first = (f"deviates at step {group(r['first_deviation'])}"
                 if r["first_deviation"] else
                 f"exact for {group(d['steps'])} steps")
        rows.append(f"{sci(r['epsilon'])} & {sci(r['gamma'])} & "
                    f"${r['worst_bound']:.3f}$ & "
                    f"{'yes' if r['certified_by_corollary'] else 'no'} & "
                    f"{first}\\\\")
    body = "\n".join(rows)
    return rf"""\begin{{table}}[ht]
\centering\small
\begin{{tabular}}{{@{{}}lllll@{{}}}}
\toprule
$\varepsilon$ & $\gamma$ & $\max_{{\ell,i}}(\varepsilon k+\gamma(n-k))$ &
Certified & Outcome\\
\midrule
{body}
\bottomrule
\end{{tabular}}
\caption{{Static weight error on $\Lev(\Nhost)$: each programmed weight is
perturbed uniformly in $[-\varepsilon,\varepsilon]$ and each zero entry
uniformly in $[-\gamma,\gamma]$, and the perturbed map is run beside the exact
one from the initial state of the self-describing instance until their states
differ. The widest layer has $n_{{\max}}={group(d['widest_layer'])}$ and the
largest effective fan-in is {d['largest_fanin']}, so
Corollary~\ref{{cor:mismatch}} certifies exactly the rows whose third column is
below $\tfrac12$.}}
\label{{tab:leak}}
\end{{table}}"""


def curve_table():
    d = load("paper_noise_curve.json")
    rows = []
    for r in d["rows"]:
        rate = (sci(r["empirical_rate_per_step"])
                if r["empirical_rate_per_step"] else "none seen")
        ratio = f"${r['looseness']:.1f}$" if r["looseness"] else "---"
        med = (group(r["median_first_deviation"])
               if r["median_first_deviation"] else "---")
        rows.append(f"${r['sigma']:.2f}$ & {r['deviated']}/{r['copies']} & "
                    f"{med} & {group(r['step_exposure'])} & {rate} & "
                    f"{sci(r['bound_rate_per_step'])} & {ratio}\\\\")
    body = "\n".join(rows)
    return rf"""\begin{{table}}[ht]
\centering\small
\begin{{tabular}}{{@{{}}lllrllr@{{}}}}
\toprule
$s$ & Deviated & Median step & Exposure & Measured rate & Bound $Mp(s)$ &
Ratio\\
\midrule
{body}
\bottomrule
\end{{tabular}}
\caption{{The probability that one step of $\Lev(\Nhost)$ deviates from the
exact orbit under read noise of standard deviation $s$. For each $s$, one
noise-free copy and {d['rows'][0]['copies']} noisy copies were stepped together
from the initial state of the self-describing instance and the full state
compared after each step; the exposure is the number of step--copy pairs before
a first deviation. The bound is that of Proposition~\ref{{prop:random}} with
$M={group(d['M'])}$, and the last column is the ratio of the two. Below
$s=0.08$ the bound is smaller than any rate this experiment could resolve.}}
\label{{tab:curve}}
\end{{table}}"""


def speed_table():
    d = load("paper_throughput.json")
    T = json.load(open(os.path.join(REPO, "paper", "runs",
                                    "paper_runs_reference.json"))
                  )["selfrep"]["generations"][0]["steps"]
    order = [("reference", "integer reference, Definition 3.1"),
             ("net_cpu", "the netlist of $\\sigma$, unit by unit, one core"),
             ("net_gpu", "the netlist of $\\sigma$, unit by unit, GPU"),
             ("dense_gpu", "$\\Lev(\\Nhost)$ as dense matrices, GPU"),
             ("graph_gpu", "$\\Lev(\\Nhost)$ from its nonzero entries, GPU"),
             ("sparse_cpu", "$\\Lev(\\Nhost)$ from its nonzero entries, one core")]
    rows = []
    for key, what in order:
        if key not in d:
            continue
        r = d[key]["steps_per_second"]
        gen = T / r
        span = (f"{gen:.1f} s" if gen < 90 else f"{gen / 60:.1f} min")
        rows.append(f"{what} & {group(round(r))} & {span}\\\\")
    body = "\n".join(rows)
    ent = d.get("dense_gpu", {}).get("entries", 0)
    nz = d.get("dense_gpu", {}).get("nonzero", 0)
    bw = d.get("dense_gpu", {}).get("bandwidth_gbs", 0)
    return rf"""\begin{{table}}[ht]
\centering\small
\begin{{tabular}}{{@{{}}lrr@{{}}}}
\toprule
Evaluator & Steps per second & One generation\\
\midrule
{body}
\bottomrule
\end{{tabular}}
\caption{{What one generation of $\Omega^\ast$ costs on each evaluator, measured
on an otherwise idle machine over $20{{,}}000$ steps of the self-describing
instance after a warm-up. The dense form reads all
{group(ent)} matrix entries at every step, of which {group(nz)} are nonzero,
which is {ent * 4 / 1e6:.0f}~MB and {bw:.0f}~GB/s, close to the memory bandwidth
of the accelerator. Evaluating the same map from its nonzero entries replaces
that traffic by {group(nz)} gathers and is four times faster once the level plan
is captured as a single CUDA graph, which issues one operation per step in place
of some hundreds; on one processor core, where there is no launch to amortise,
the gathers cost more than the dense arithmetic saves. The netlist carries
$50{{,}}250$ weighted predecessor entries against the {group(nz)} of its
levelization and is faster in that proportion on either processor.}}
\label{{tab:speed}}
\end{{table}}"""


def main() -> int:
    if len(sys.argv) > 1 and sys.argv[1] == "speed":
        print(speed_table())
        return 0
    if len(sys.argv) > 1 and sys.argv[1] == "tables":
        print(omega_table())
        print()
        print(curve_table())
        print()
        print(leak_table())
        return 0
    d = json.load(open(os.path.join(REPO, "paper", "runs",
                                    "paper_noise_curve.json")))
    rows = [r for r in d["rows"]]
    xs = [r["sigma"] for r in rows]
    emp = [(r["sigma"], r["empirical_rate_per_step"]) for r in rows
           if r["empirical_rate_per_step"] > 0]
    bnd = [(r["sigma"], r["bound_rate_per_step"]) for r in rows]

    lo = math.floor(min([math.log10(y) for _, y in emp]
                        + [math.log10(y) for _, y in bnd if y > 0]))
    hi = math.ceil(max([math.log10(y) for _, y in emp]
                       + [math.log10(y) for _, y in bnd]))
    x0, x1 = min(xs), max(xs)
    W, H = 105.0, 55.0

    def px(s):
        return W * (s - x0) / (x1 - x0)

    def py(v):
        return H * (math.log10(v) - lo) / (hi - lo)

    out = []
    a = out.append
    a(r"\begin{figure}[ht]")
    a(r"\centering")
    a(r"\begin{tikzpicture}[font=\small, x=1mm, y=1mm]")
    a(rf"\draw[->,>=stealth] (-2,0) -- ({W + 6:.1f},0) "
      rf"node[right,font=\footnotesize] {{$s$}};")
    a(rf"\draw[->,>=stealth] (-2,-2) -- (-2,{H + 6:.1f});")
    for e in range(lo, hi + 1):
        y = H * (e - lo) / (hi - lo)
        a(rf"\draw[black!20] (-2,{y:.2f}) -- ({W:.2f},{y:.2f});")
        a(rf"\node[left,font=\scriptsize] at (-2.5,{y:.2f}) "
          rf"{{$10^{{{e}}}$}};")
    for s in xs:
        a(rf"\draw (({px(s):.2f}),0) -- ({px(s):.2f},-1.2);")
        a(rf"\node[below,font=\scriptsize] at ({px(s):.2f},-1.2) "
          rf"{{${s:g}$}};")
    a(rf"\node[rotate=90,font=\footnotesize] at (-12,{H / 2:.1f}) "
      rf"{{probability per step}};")
    pts = " -- ".join(f"({px(s):.2f},{py(v):.2f})" for s, v in bnd)
    a(rf"\draw[thick,dashed] {pts};")
    for s, v in bnd:
        a(rf"\fill ({px(s):.2f},{py(v):.2f}) circle (0.7mm);")
    pts = " -- ".join(f"({px(s):.2f},{py(v):.2f})" for s, v in emp)
    a(rf"\draw[thick] {pts};")
    for s, v in emp:
        a(rf"\draw[fill=white,thick] ({px(s):.2f},{py(v):.2f}) "
          rf"circle (0.8mm);")
    a(rf"\draw[dashed,black!50] ({px(x0):.2f},{py(1.0):.2f}) -- "
      rf"({W:.2f},{py(1.0):.2f});")
    a(rf"\node[right,font=\scriptsize] at ({W - 30:.1f},{py(1.0) + 3:.2f}) "
      rf"{{bound vacuous above here}};")
    a(rf"\node[font=\scriptsize,anchor=west] at (4,{py(bnd[-1][1]) - 5:.2f}) "
      rf"{{union bound $Mp(s)$}};")
    a(rf"\node[font=\scriptsize,anchor=west] at (20,{py(emp[0][1]) - 5:.2f}) "
      rf"{{measured}};")
    a(r"\end{tikzpicture}")
    a(r"\caption{The probability that one step of $\Lev(\Nhost)$ deviates from "
      r"the exact orbit under additive Gaussian read noise of standard "
      r"deviation $s$, measured by running noise-free and noisy copies together "
      r"and comparing the full state after every step, against the union bound "
      r"$Mp(s)$ of Proposition~\ref{prop:random} with $M=87{,}294$. Both curves "
      r"are the same quantity, the upper one as bounded above.}")
    a(r"\label{fig:curve}")
    a(r"\end{figure}")
    print("\n".join(out))
    return 0


if __name__ == "__main__":
    sys.exit(main())