#!/usr/bin/env python3
# =============================================================================
# figures.py — paper figures from existing campaign data
# Author: Simon-Pierre Boucher — contact@spboucher.ai
# =============================================================================
# Generates (into paper/figs/):
#   fig_scaling.pdf     rho·lnx vs lnx: data, fit, confirmed predictions, HL model
#   fig_race.pdf        D(x) = N2-N4 lead changes (fine trace to 1e10 + checkpoints)
#   fig_hist.pdf        N(g, 1e8) with the mod-6 comb structure
#   fig_certificate.pdf the 260-gap desert certificate map
# Palette (validated, dataviz method): blue #2a78d6, orange #eb6834, aqua #1baf7a
# Deterministic; the race trace is recomputed locally (scan to 1e10, ~40 s).
# =============================================================================

import json
from math import log
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

from core import primes_upto, sieve_segment

ROOT = Path(__file__).resolve().parent.parent
FIGS = ROOT / "paper" / "figs"
FIGS.mkdir(exist_ok=True)
DATA = ROOT / "data"

BLUE, ORANGE, AQUA = "#2a78d6", "#eb6834", "#1baf7a"
INK, MUTED = "#0b0b0b", "#52514e"

plt.rcParams.update({
    "font.size": 8.5, "axes.labelsize": 9, "axes.titlesize": 9.5,
    "axes.spines.top": False, "axes.spines.right": False,
    "axes.linewidth": 0.7, "xtick.major.width": 0.7, "ytick.major.width": 0.7,
    "grid.linewidth": 0.5, "grid.alpha": 0.25, "legend.frameon": False,
    "figure.dpi": 150, "text.color": INK, "axes.edgecolor": MUTED,
    "xtick.color": MUTED, "ytick.color": MUTED, "axes.labelcolor": INK,
})


def load_checkpoints():
    c2 = json.load(open(DATA / "cycle2_c4_4e10_M3U96a.json"))["checkpoints"]
    c4 = json.load(open(DATA / "cycle4_c4_1e13.json"))["checkpoints"]
    pts = [(float(x), v["rho1_lnx"], v["rho2_lnx"], v["D_N2_minus_N4"])
           for x, v in c2.items() if 1e8 <= float(x) < 1e10]
    pts += [(float(x), v["rho1_lnx"], v["rho2_lnx"], v["D_N2_minus_N4"])
            for x, v in c4.items()]
    return sorted(pts)


def fit_cd(lams, vals):
    A = np.vstack([np.ones_like(lams), 1 / lams]).T
    (c, d), *_ = np.linalg.lstsq(A, vals, rcond=None)
    return c, d


def fig_scaling(pts):
    lam = np.array([log(x) for x, *_ in pts])
    v1 = np.array([p[1] for p in pts])
    v2 = np.array([p[2] for p in pts])
    c1, d1 = fit_cd(lam, v1)
    c2_, d2_ = fit_cd(lam, v2)
    model = json.load(open(DATA / "model_c4.json"))["predictions"]
    ml = np.array([m["lambda"] for m in model.values()])
    mv = np.array([m["rho_times_lambda"] for m in model.values()])

    fig, (ax, axb) = plt.subplots(1, 2, figsize=(6.4, 2.7),
                                  gridspec_kw={"width_ratios": [3, 2]})
    ll = np.linspace(18.0, 34.6, 200)
    # panel (a): data + fits + confirmed predictions
    ax.plot(ll, c1 + d1 / ll, color=BLUE, lw=1.4)
    ax.plot(ll, c2_ + d2_ / ll, color=AQUA, lw=1.4)
    ax.plot(lam, v1, "o", ms=3.4, color=BLUE, mfc="white", mew=1.1)
    ax.plot(lam, v2, "o", ms=3.4, color=AQUA, mfc="white", mew=1.1)
    # confirmed out-of-sample predictions (1e12 and 1e13 for lag1; 1e13 for lag2)
    for X, obs in ((1e12, -0.54756), (1e13, -0.54264)):
        ax.plot(log(X), obs, "D", ms=5.5, color=ORANGE, mfc="none", mew=1.4)
    ax.plot(log(1e13), -0.25247, "D", ms=5.5, color=ORANGE, mfc="none", mew=1.4)
    for X, txt in ((1e14, -0.5387), (1e15, -0.5350)):
        ax.plot(log(X), txt, "s", ms=4.5, color=MUTED, mfc="none", mew=1.0)
    ax.annotate("lag 1:  $\\rho\\,\\ln x = c + d/\\ln x$", (18.4, -0.525),
                color=BLUE, fontsize=8)
    ax.annotate("lag 2", (19.8, -0.272), color=AQUA, fontsize=8)
    ax.annotate("confirmed predictions", (29.2, -0.572),
                color=ORANGE, fontsize=7.5, ha="center")
    ax.annotate("", xy=(log(1e12) + 0.1, -0.553), xytext=(28.2, -0.568),
                arrowprops=dict(arrowstyle="->", color=ORANGE, lw=0.9))
    ax.annotate("", xy=(log(1e13) - 0.1, -0.547), xytext=(30.0, -0.568),
                arrowprops=dict(arrowstyle="->", color=ORANGE, lw=0.9))
    ax.annotate("next targets", (31.6, -0.528), color=MUTED, fontsize=7.5)
    ax.set_xlabel("$\\ln x$")
    ax.set_ylabel("$\\rho(x)\\,\\ln x$")
    ax.grid(True)
    ax.set_xlim(17.8, 35.4)
    for X in (1e8, 1e10, 1e12, 1e14):
        ax.axvline(log(X), color=MUTED, lw=0.4, alpha=0.18)
        ax.text(log(X), ax.get_ylim()[1], f"$10^{{{int(round(log(X,10)))}}}$",
                fontsize=6.5, color=MUTED, ha="center", va="bottom")
    # panel (b): observed vs first-order HL model
    axb.plot(lam, v1, "o", ms=3.2, color=BLUE, mfc="white", mew=1.0, label="observed (lag 1)")
    axb.plot(ll, c1 + d1 / ll, color=BLUE, lw=1.2)
    axb.plot(ml, mv, "--", color=ORANGE, lw=1.4, label="first-order HL model")
    axb.axhline(0, color=MUTED, lw=0.6)
    axb.set_xlim(17.8, 95)
    axb.set_ylim(-0.65, 0.02)
    axb.set_xlabel("$\\lambda = \\ln x$")
    axb.grid(True)
    axb.annotate("observed", (23, -0.52), color=BLUE, fontsize=8)
    axb.annotate("first-order HL model\n($\\approx 32\\%$ of magnitude)",
                 (40, -0.115), color=ORANGE, fontsize=8)
    fig.tight_layout()
    fig.savefig(FIGS / "fig_scaling.pdf")
    plt.close(fig)
    print("fig_scaling.pdf")


def race_trace(limit=10**10, per_decade=60):
    """D(x) sampled at ~per_decade log-spaced points per decade, exact."""
    targets = np.unique(np.round(np.logspace(6, np.log10(limit),
                        int(per_decade * (np.log10(limit) - 6)))).astype(np.int64))
    base = primes_upto(int(limit ** 0.5) + 1)
    D = 0
    out_x, out_d = [], []
    prev = None
    ti = 0
    for lo in range(2, limit, 50_000_000):
        hi = min(lo + 50_000_000, limit)
        primes = sieve_segment(lo, hi, base)
        if prev is not None:
            primes = np.concatenate(([prev], primes))
        gaps = np.diff(primes)
        ends = primes[1:]
        delta = (gaps == 2).astype(np.int64) - (gaps == 4).astype(np.int64)
        run = D + np.cumsum(delta)
        while ti < len(targets) and targets[ti] < hi:
            j = np.searchsorted(ends, targets[ti], side="right") - 1
            out_x.append(int(targets[ti]))
            out_d.append(int(run[j]) if j >= 0 else D)
            ti += 1
        D = int(run[-1])
        prev = primes[-1]
    return np.array(out_x, dtype=float), np.array(out_d, dtype=float)


def fig_race(pts):
    xs, ds = race_trace()
    fig, ax = plt.subplots(figsize=(6.4, 2.5))
    ax.axhline(0, color=MUTED, lw=0.7)
    ax.plot(xs, ds, color=BLUE, lw=1.1)
    cx = [p[0] for p in pts if p[0] >= 2e10]
    cd = [p[3] for p in pts if p[0] >= 2e10]
    ax.plot(cx, cd, "D", ms=4.5, color=ORANGE, mfc="none", mew=1.2)
    ax.set_xscale("log")
    ax.set_yscale("symlog", linthresh=200)
    ax.set_xlim(1e6, 2e13)
    ax.set_ylim(-2.5e5, 6e4)
    ax.set_xlabel("$x$")
    ax.set_ylabel("$D(x) = N(2,x) - N(4,x)$")
    ax.grid(True, which="major")
    ax.annotate("first tie after $10^6$\n$x = 80{,}966{,}861$",
                xy=(8.1e7, 0), xytext=(2.5e6, -2600), fontsize=7.5, color=INK,
                arrowprops=dict(arrowstyle="->", color=MUTED, lw=0.8))
    ax.annotate("fine trace (every gap, sampled)", (2.6e6, 6000), color=BLUE, fontsize=8)
    ax.annotate("campaign checkpoints", (2.2e11, -12000), color=ORANGE, fontsize=8)
    fig.tight_layout()
    fig.savefig(FIGS / "fig_race.pdf")
    plt.close(fig)
    print("fig_race.pdf")


def fig_hist():
    rows = np.loadtxt(DATA / "gap_histogram_1e8.csv", delimiter=",", skiprows=1)
    g, cnt = rows[:, 0].astype(int), rows[:, 1]
    m = (g >= 2) & (g <= 150) & (g % 2 == 0)
    g, cnt = g[m], cnt[m]
    six = g % 6 == 0
    fig, ax = plt.subplots(figsize=(6.4, 2.5))
    ax.vlines(g[~six], 1, cnt[~six], color=MUTED, lw=1.6, alpha=0.75)
    ax.vlines(g[six], 1, cnt[six], color=BLUE, lw=1.9)
    ax.set_yscale("log")
    ax.set_xlabel("gap $g$")
    ax.set_ylabel("$N(g, 10^8)$")
    ax.grid(True, axis="y")
    ax.annotate("multiples of 6", (66, 2.2e5), color=BLUE, fontsize=8.5)
    ax.annotate("other even gaps", (80, 1.1e4), color=MUTED, fontsize=8.5)
    ax.annotate("jumping champion $g = 6$", xy=(6, 8.8e5), xytext=(16, 1.5e6),
                fontsize=7.5, color=INK,
                arrowprops=dict(arrowstyle="->", color=MUTED, lw=0.8))
    ax.set_xlim(0, 152)
    ax.set_ylim(1, 4e6)
    fig.tight_layout()
    fig.savefig(FIGS / "fig_hist.pdf")
    plt.close(fig)
    print("fig_hist.pdf")


def fig_certificate():
    cert = json.load(open(ROOT / "certs" / "desert_certificate.json"))
    interior = cert["interior_certificates"]
    cov_x, cov_y, tf_x, tf_y, mr_x = [], [], [], [], []
    for i_str, c in interior.items():
        i = int(i_str)
        if c["type"] == "covering_factor":
            cov_x.append(i); cov_y.append(c["q"])
        elif c["type"] == "trial_factor":
            tf_x.append(i); tf_y.append(c["q"])
        else:
            mr_x.append(i)
    fig, ax = plt.subplots(figsize=(6.4, 2.6))
    ax.plot(cov_x, cov_y, "o", ms=2.6, color=BLUE, mew=0)
    ax.plot(tf_x, tf_y, "D", ms=4.5, color=ORANGE, mfc="none", mew=1.2)
    top = max(tf_y) * 3
    ax.plot(mr_x, [top] * len(mr_x), "s", ms=5, color=AQUA, mfc="none", mew=1.3)
    ax.set_yscale("log")
    ax.set_xlabel("position $i$ in the desert  ($N + i$, $1 \\leq i \\leq 259$)")
    ax.set_ylabel("certifying prime factor $q_i$")
    ax.grid(True, axis="y")
    ax.annotate("covering system (proven for every shift $t$)", (4, 230),
                color=BLUE, fontsize=8)
    ax.annotate("holes: explicit trial factor", (138, 2.6e4),
                color=ORANGE, fontsize=8)
    ax.annotate("hole: strong MR witness", (mr_x[0] + 8, top * 0.75),
                color=AQUA, fontsize=8)
    ax.set_xlim(0, 260)
    fig.tight_layout()
    fig.savefig(FIGS / "fig_certificate.pdf")
    plt.close(fig)
    print("fig_certificate.pdf")


if __name__ == "__main__":
    pts = load_checkpoints()
    fig_scaling(pts)
    fig_hist()
    fig_certificate()
    fig_race(pts)   # last: recomputes a 1e10 scan (~40 s)