"""Stage 6: build the manuscript figures as PDF files in results/figures/.

matplotlib is not available in the sandbox, so every figure is emitted as a
TikZ/pgfplots document (results/figures/src/*.tex) and compiled with the
locally cached tectonic bundle.  All plotted numbers are read from results/.
Graph drawings are produced from the P(n,k) definition; chronologies are
re-validated here by an independent closure replay before they are drawn
(an illegal force aborts the build).
"""
import csv
import json
import math
import os
import subprocess
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
os.chdir(ROOT)
OUT = "results/figures"
SRC = f"{OUT}/src"
os.makedirs(SRC, exist_ok=True)
import shutil
shutil.copy("manuscript/autoscholar.cls", f"{SRC}/autoscholar.cls")  # same class and fonts as the manuscript

PRE = r"""\documentclass{autoscholar}
\usepackage{tikz}
\usepackage{pgfplots}
\pgfplotsset{compat=1.18}
\usetikzlibrary{arrows.meta,calc,backgrounds,positioning}
\definecolor{asblue}{HTML}{2B5C8A}
\definecolor{asgold}{HTML}{B8860B}
% the figure is measured in a box and shipped out on a page of exactly its size;
% the class's end-of-document disclosure is not wanted on a figure file
\RemoveFromHook{enddocument}[autoscholar]
\newsavebox\figbox
\begin{document}
\begin{lrbox}{\figbox}%
"""
POST = r"""%
\end{lrbox}
\hoffset=-1in \voffset=-1in
\pdfpagewidth=\dimexpr\wd\figbox+6pt\relax
\pdfpageheight=\dimexpr\ht\figbox+\dp\figbox+6pt\relax
\shipout\vbox{\kern3pt\hbox{\kern3pt\usebox\figbox\kern3pt}\kern3pt}
\end{document}
"""


# ------------------------------------------------------------------ graph helpers
def nbrs(n, k, v):
    layer, i = v
    if layer == "u":
        return [("u", (i - 1) % n), ("u", (i + 1) % n), ("v", i)]
    return [("v", (i - k) % n), ("v", (i + k) % n), ("u", i)]


def replay(n, k, S, chron):
    """independent legality check of a chronology; returns final black set"""
    black = set(S)
    for src, tgt in chron:
        if src not in black or tgt in black:
            raise ValueError(f"illegal force {src}->{tgt}: source white or target black")
        white = [w for w in nbrs(n, k, src) if w not in black]
        if white != [tgt]:
            raise ValueError(f"illegal force {src}->{tgt}: white neighbours {white}")
        black.add(tgt)
    return black


def parse(label):
    return (label[0], int(label[1:]))


def pos(n, layer, i, R=3.0, r=1.75):
    ang = math.pi / 2 - 2 * math.pi * i / n
    rad = R if layer == "u" else r
    return rad * math.cos(ang), rad * math.sin(ang)


def petersen_tikz(n, k, fill=None, labels=None, forces=None, R=3.0, r=1.75, node_r="2.6pt", label_size="tiny", show_index=True, legend=None):
    """fill: dict vertex -> colour ; labels: dict vertex -> text inside node;
    forces: list of (src, tgt, step) drawn as arrows along the edge"""
    fill = fill or {}
    labels = labels or {}
    L = []
    L.append("\\begin{tikzpicture}[every node/.style={inner sep=0pt}]")
    L.append(f"\\tikzset{{v/.style={{circle,draw=asink,line width=0.5pt,minimum size={node_r},inner sep=0pt}}}}")
    # edges
    for i in range(n):
        x1, y1 = pos(n, "u", i, R, r)
        x2, y2 = pos(n, "u", (i + 1) % n, R, r)
        L.append(f"\\draw[asstone,line width=0.6pt] ({x1:.3f},{y1:.3f}) -- ({x2:.3f},{y2:.3f});")
        x3, y3 = pos(n, "v", i, R, r)
        L.append(f"\\draw[asline,line width=0.6pt] ({x1:.3f},{y1:.3f}) -- ({x3:.3f},{y3:.3f});")
        x4, y4 = pos(n, "v", (i + k) % n, R, r)
        L.append(f"\\draw[asstone,line width=0.6pt] ({x3:.3f},{y3:.3f}) -- ({x4:.3f},{y4:.3f});")
    # forces
    for src, tgt, step in (forces or []):
        xs, ys = pos(n, src[0], src[1], R, r)
        xt, yt = pos(n, tgt[0], tgt[1], R, r)
        L.append(f"\\draw[asaccent,line width=0.9pt,-{{Stealth[length=4pt]}},shorten >=3.2pt,shorten <=3.2pt] ({xs:.3f},{ys:.3f}) -- ({xt:.3f},{yt:.3f});")
    # vertices
    for layer in ("u", "v"):
        for i in range(n):
            v = (layer, i)
            x, y = pos(n, layer, i, R, r)
            col = fill.get(v, "white")
            txt = labels.get(v, "")
            tcol = "white" if col in ("asink", "asaccent", "asblue") else "asink"
            L.append(f"\\node[v,fill={col}] at ({x:.3f},{y:.3f}) {{\\{label_size}\\color{{{tcol}}}{txt}}};")
            if show_index:
                ang = math.pi / 2 - 2 * math.pi * i / n
                off = 0.36 if layer == "u" else -0.37
                lx, ly = (R + off) * math.cos(ang), (R + off) * math.sin(ang)
                if layer == "v":
                    lx, ly = (r + off) * math.cos(ang), (r + off) * math.sin(ang)
                # index labels at \scriptsize (enlarged after referee round 1, finding F7)
                L.append(f"\\node[inner sep=1pt] at ({lx:.3f},{ly:.3f}) {{\\scriptsize\\color{{asstone}}${layer}_{{{i}}}$}};")
    if legend:
        L.append(legend)
    L.append("\\end{tikzpicture}")
    return "\n".join(L)


def compile_tex(name, body):
    path = f"{SRC}/{name}.tex"
    open(path, "w").write(PRE + body + POST)
    cmd = ["tectonic", "-X", "compile", "--only-cached", "-o", OUT, path]
    res = subprocess.run(cmd, capture_output=True, text=True)
    if res.returncode != 0:
        sys.stderr.write(res.stdout[-3000:] + res.stderr[-3000:])
        raise SystemExit(f"figure {name} failed")
    print("built", f"{OUT}/{name}.pdf")


# ------------------------------------------------------------------ data
exact = {int(r["n"]): r for r in csv.DictReader(open("results/exact_values.csv"))}
kvn = {int(r["n"]): r for r in csv.DictReader(open("results/known_vs_new.csv"))}
window = list(csv.DictReader(open("results/tables/boundary_window_p18.csv")))
btab = list(csv.DictReader(open("results/tables/boundary_table.csv")))
bf = list(csv.DictReader(open("results/tables/brute_force.csv")))
wit = [json.loads(l) for l in open("results/upper_bound_witnesses.jsonl") if l.strip()]
cex = [json.loads(l) for l in open("results/counterexamples.jsonl") if l.strip()]
summary = json.load(open("results/robustness/summary.json"))

# ------------------------------------------------------------------ Fig 1: anatomy P(13,5) and P(15,5)
for n in (13, 15):
    body = petersen_tikz(n, 5, R=2.6, r=1.5, node_r="3.4pt")
    compile_tex(f"fig1_anatomy_P{n}_5", body)

# ------------------------------------------------------------------ Fig 2a: P7 chronology on P(14,5)
n = 14
S = [("u", i) for i in range(12)]
chron = [(("u", i), ("v", i)) for i in range(1, 11)] + [(("v", 5), ("v", 0)), (("v", 6), ("v", 11))]
for j in range(12, n):
    chron += [(("u", j - 1), ("u", j)), (("v", j - 5), ("v", j))]
final = replay(n, 5, S, chron)
assert len(final) == 2 * n, "P7 chronology does not force P(14,5)"
fill = {v: "asink" for v in S}
labels = {}
forces = []
for step, (s, t) in enumerate(chron, 1):
    labels[t] = str(step)
    forces.append((s, t, step))
legend = ("\\node[anchor=north,align=left,text width=7.6cm] at (0,-3.55) {\\footnotesize black $=$ initial set $S=\\{u_0,\\dots,u_{11}\\}$;\\ "
          "numbers $=$ step at which the vertex is forced};")
compile_tex("fig2a_chronology_P14_5_P7", petersen_tikz(n, 5, fill=fill, labels=labels, forces=forces, R=2.9, r=1.7, node_r="9pt", label_size="scriptsize", legend=legend))

# ------------------------------------------------------------------ Fig 2b: optimal witness on P(15,5)
w15 = min((w for w in wit if w["n"] == 15), key=lambda w: w["size"])
S = [parse(x) for x in w15["S"]]
chron = [(parse(a), parse(b)) for a, b in w15["chronology"]]
final = replay(15, 5, S, chron)
assert len(final) == 30 and len(S) == 8
fill = {v: "asink" for v in S}
labels = {}
forces = []
for step, (s, t) in enumerate(chron, 1):
    labels[t] = str(step)
    forces.append((s, t, step))
legend = ("\\node[anchor=north,align=left,text width=7.6cm] at (0,-3.55) {\\footnotesize black $=$ minimum zero-forcing set of size $8$ "
          "($" + ",".join(f"{a}_{{{i}}}" for a, i in S) + "$)};")
compile_tex("fig2b_chronology_P15_5_optimal", petersen_tikz(15, 5, fill=fill, labels=labels, forces=forces, R=2.9, r=1.7, node_r="9pt", label_size="scriptsize", legend=legend))

# ------------------------------------------------------------------ Fig 3: exact values 11..48
known = [(n, int(exact[n]["Z"])) for n in sorted(exact) if kvn[n]["classification"] == "known"]
new = [(n, int(exact[n]["Z"])) for n in sorted(exact) if kvn[n]["classification"] == "new"]
hold = [(n, int(exact[n]["Z"])) for n in sorted(exact) if exact[n]["holdout"] == "True"]
ub = [(n, int(kvn[n]["published_upper_bound"])) for n in sorted(exact)]
body = r"""\begin{tikzpicture}
\begin{axis}[width=14cm,height=7cm,xlabel={$n$},ylabel={$Z(P(n,5))$},xmin=10,xmax=49,ymin=2,ymax=13.5,
  ytick={2,4,6,8,10,12},xtick={11,15,20,25,30,35,40,45,48},grid=major,grid style={asline!60},
  legend style={at={(0.98,0.05)},anchor=south east,font=\footnotesize,draw=asline,fill=white,row sep=-1pt},
  axis line style={asstone},tick label style={font=\footnotesize},label style={font=\small},
  every axis plot/.append style={line width=0.7pt}]
\addplot[asstone,dashed,line width=0.6pt] coordinates {(10,12) (49,12)};
\addlegendentry{published upper bound $12$ (all $n\ge 11$)}
\addplot[asstone,densely dotted,line width=0.6pt] coordinates {(15,6) (49,6)};
\addlegendentry{published lower bound $6$ ($n\ge 15$)}
\addplot[asline,only marks,mark=triangle*,mark size=2.6pt,mark options={fill=asline,draw=asstone}] coordinates {""" + " ".join(f"({n},{z})" for n, z in ub if n <= 22) + r"""};
\addlegendentry{best published upper bound, $n\le 22$}
\addplot[asblue,only marks,mark=square*,mark size=2.6pt] coordinates {""" + " ".join(f"({n},{z})" for n, z in known) + r"""};
\addlegendentry{exact value known before (see text)}
\addplot[asaccent,only marks,mark=*,mark size=2.4pt] coordinates {""" + " ".join(f"({n},{z})" for n, z in new) + r"""};
\addlegendentry{exact value certified here (new to our knowledge)}
\addplot[asgold,only marks,mark=o,mark size=4.2pt,line width=0.9pt] coordinates {""" + " ".join(f"({n},{z})" for n, z in hold) + r"""};
\addlegendentry{pre-registered holdout $n$}
\draw[asaccent,line width=0.5pt] (axis cs:22.5,2) -- (axis cs:22.5,13.5);
\node[anchor=north west,font=\footnotesize,text=asaccent] at (axis cs:22.7,13.4) {$n_0=23$: $Z=12$ for all $n\ge 23$ (Theorem~\ref{thm:main})};
\end{axis}
\end{tikzpicture}"""
compile_tex("fig3_exact_values", body.replace("(Theorem~\\ref{thm:main})", "(Theorem 1.1)"))

# ------------------------------------------------------------------ Fig 4: obstruction at n = 25, p = 18
ob = next(c for c in cex if c.get("type") == "boundary_obstruction" and c["n"] == 25 and c["p"] == 18)
X = [parse(x) for x in ob["X"]]
B = [parse(x) for x in ob["boundary"]]
# independent recount of the boundary
Bchk = set()
for v in X:
    for w in nbrs(25, 5, v):
        if w not in set(X):
            Bchk.add(w)
assert Bchk == set(B) and len(B) == 10 and len(X) == 18, "obstruction recount failed"
fill = {v: "asaccent" for v in X}
fill.update({v: "asstone" for v in B})
legend = ("\\node[anchor=north,align=left,text width=7.6cm] at (0,-3.6) {\\footnotesize rust $=X$ ($|X|=18$: the two inner pentagons $\\{v_i: i\\equiv 0,1 \\pmod 5\\}$ and the eight outer vertices $u_i$, $i\\equiv 0,1 \\pmod 5$, $i\\ne 0,21$);\\ "
          "grey $=\\partial X$ ($|\\partial X|=10$); every other vertex has no neighbour in $X$};")
compile_tex("fig4_obstruction_P25_5", petersen_tikz(25, 5, fill=fill, R=3.0, r=1.75, node_r="5pt", show_index=False, legend=legend))

# ------------------------------------------------------------------ Fig 5a: boundary window scaling
w = [(int(r["n"]), float(r["solve_seconds"]), float(r["check_seconds"]), int(r["proof_lines"])) for r in window if int(r["n"]) >= 26]
body = r"""\begin{tikzpicture}
\begin{axis}[width=13.5cm,height=6.5cm,xlabel={$n$ (certified window $26\le n\le 154$, $p=18$)},ylabel={seconds},
  xmin=24,xmax=156,ymin=0,ymax=290,grid=major,grid style={asline!60},axis line style={asstone},
  tick label style={font=\footnotesize},label style={font=\small},
  legend style={at={(0.02,0.97)},anchor=north west,font=\footnotesize,draw=asline,fill=white},
  axis y line*=left]
\addplot[asaccent,only marks,mark=*,mark size=1.3pt] coordinates {""" + " ".join(f"({n},{s})" for n, s, c, p in w) + r"""};
\addlegendentry{CaDiCaL 1.9.5 solve time (UNSAT, DRAT logging)}
\addplot[asblue,only marks,mark=square*,mark size=1.3pt] coordinates {""" + " ".join(f"({n},{c})" for n, s, c, p in w) + r"""};
\addlegendentry{drat-trim verification time}
\end{axis}
\begin{axis}[width=13.5cm,height=6.5cm,xmin=24,xmax=156,ymin=0,ymax=5.5,axis y line*=right,axis x line=none,
  ylabel={DRAT proof lines ($10^6$)},tick label style={font=\footnotesize},label style={font=\small},
  legend style={at={(0.02,0.80)},anchor=north west,font=\footnotesize,draw=asline,fill=white},axis line style={asstone}]
\addplot[asstone,only marks,mark=triangle*,mark size=1.4pt] coordinates {""" + " ".join(f"({n},{p/1e6:.3f})" for n, s, c, p in w) + r"""};
\addlegendentry{proof length}
\end{axis}
\end{tikzpicture}"""
compile_tex("fig5a_window_scaling", body)

# ------------------------------------------------------------------ Fig 5b: exhaustive enumeration cost
Z = {n: int(exact[n]["Z"]) for n in exact}
bfrows = [(int(r["n"]), int(r["s"]), float(r["wall_seconds"]), int(r["subsets"])) for r in bf]
pts = [(n, t, sub) for n, s, t, sub in bfrows if n >= 17 and s == Z[n] - 1]
body = r"""\begin{tikzpicture}
\begin{semilogyaxis}[width=13.5cm,height=6.5cm,xlabel={$n$},ylabel={wall-clock seconds (16 threads)},xmin=16,xmax=28,ymin=0.1,ymax=3000,
  grid=major,grid style={asline!60},axis line style={asstone},tick label style={font=\footnotesize},label style={font=\small},
  legend style={at={(0.02,0.97)},anchor=north west,font=\footnotesize,draw=asline,fill=white},xtick={17,...,27}]
\addplot[asaccent,mark=*,mark size=1.8pt,line width=0.7pt] coordinates {""" + " ".join(f"({n},{t})" for n, t, sub in pts) + r"""};
\addlegendentry{$D_n$-orbit enumeration of all $(Z-1)$-subsets (C oracle)}
\end{semilogyaxis}
\begin{semilogyaxis}[width=13.5cm,height=6.5cm,xmin=16,xmax=28,ymin=1e6,ymax=1e12,axis y line*=right,axis x line=none,
  ylabel={subsets of size $Z-1$},tick label style={font=\footnotesize},label style={font=\small},axis line style={asstone},
  legend style={at={(0.02,0.82)},anchor=north west,font=\footnotesize,draw=asline,fill=white}]
\addplot[asstone,mark=triangle*,mark size=1.8pt,line width=0.6pt,dashed] coordinates {""" + " ".join(f"({n},{sub})" for n, t, sub in pts) + r"""};
\addlegendentry{$\binom{2n}{Z-1}$}
\end{semilogyaxis}
\end{tikzpicture}"""
compile_tex("fig5b_enumeration_scaling", body)

# ------------------------------------------------------------------ Fig 6: heat map of beta_p(n)
cols = [c for c in btab[0].keys() if c != "p"]
cw, ch = 0.62, 0.5
L = ["\\begin{tikzpicture}"]
for ri, r in enumerate(btab):
    p = int(r["p"])
    y = -ri * ch
    L.append(f"\\node[anchor=east,font=\\footnotesize] at (-0.15,{y:.2f}) {{{p}}};")
    for ci, c in enumerate(cols):
        x = ci * cw
        v = r[c]
        if v.startswith(">="):
            colr, txt, tcol = "asaccent", "$\\ge$12", "white"
        else:
            val = int(v)
            shade = int(8 + (val - 2) * 6)  # 2 -> 8%, 11 -> 62%
            colr, txt, tcol = f"asblue!{shade}", str(val), ("white" if shade > 45 else "asink")
        L.append(f"\\fill[{colr}] ({x:.2f},{y - ch/2:.2f}) rectangle ({x + cw:.2f},{y + ch/2:.2f});")
        L.append(f"\\node[font=\\footnotesize,text={tcol}] at ({x + cw/2:.2f},{y:.2f}) {{{txt}}};")
for ci, c in enumerate(cols):
    L.append(f"\\node[anchor=south,font=\\footnotesize] at ({ci*cw + cw/2:.2f},{ch/2 + 0.05:.2f}) {{{c}}};")
L.append(f"\\node[anchor=south,font=\\small] at ({len(cols)*cw/2:.2f},{ch/2 + 0.42:.2f}) {{$n$}};")
L.append(f"\\node[anchor=east,font=\\small,rotate=90] at (-0.75,{-(len(btab)-1)*ch/2:.2f}) {{$p$}};")
# highlight n = 25 column and the chosen (p=18, n>=26) region
i25 = cols.index("25")
L.append(f"\\draw[asgold,line width=1.2pt] ({i25*cw:.2f},{ch/2:.2f}) rectangle ({(i25+1)*cw:.2f},{-(len(btab)-1)*ch - ch/2:.2f});")
ip = [int(r["p"]) for r in btab].index(18)
i26 = cols.index("26")
L.append(f"\\draw[asink,line width=1.0pt,dashed] ({i26*cw:.2f},{-ip*ch + ch/2:.2f}) rectangle ({len(cols)*cw:.2f},{-ip*ch - ch/2:.2f});")
L.append(f"\\node[anchor=north,font=\\footnotesize,align=left,text width={len(cols)*cw-0.2:.1f}cm] at ({len(cols)*cw/2:.2f},{-(len(btab)-1)*ch - ch/2 - 0.15:.2f}) "
         "{cell $=\\beta_p(n)$, the minimum boundary of a $p$-set in $P(n,5)$; rust $=$ certified $\\ge 12$ (DRAT); "
         "gold frame $= n=25$, where every $p\\le 24$ fails; dashed $=$ the row $p=18$, $n\\ge 26$ used in the proof};")
L.append("\\end{tikzpicture}")
compile_tex("fig6_boundary_heatmap", "\n".join(L))

print("all figures built")