"""Regenerate the data-bearing figure (Fig. 3, Experiment 4) from the counts.

Produces `fig_exp4_core_depth.png/.pdf`: dump probability vs neutral-sector
core depth on a log scale, with 95% Wilson error bars and the non-neutral
control band.  (Figs. 1 and 2 in the paper are TikZ schematics with no
underlying data and are not regenerated here.)
"""

from __future__ import annotations

import matplotlib

matplotlib.use("Agg")  # headless backend; safe for batch runs
import matplotlib.pyplot as plt
import numpy as np

import data_io as hw          # CSV-backed loader (falls back to hardware_data)
from stats import wilson_ci


def fig_exp4(path_stem: str = "fig_exp4_core_depth") -> str:
    depths, points, err_lo, err_hi = [], [], [], []
    for r in hw.EXP4:
        if r.kind != "neutral":
            continue
        ci = wilson_ci(r.x, r.n)
        depths.append(r.extra["depth"])
        points.append(ci.point)
        err_lo.append(ci.point - ci.low)
        err_hi.append(ci.high - ci.point)

    order = sorted(range(len(depths)), key=lambda i: depths[i])
    depths = [depths[i] for i in order]
    points = [max(points[i], 1e-5) for i in order]   # clamp 0 for log axis
    err_lo = [err_lo[i] for i in order]
    err_hi = [err_hi[i] for i in order]

    control = wilson_ci(*[(r.x, r.n) for r in hw.EXP4
                          if r.label == "Control, 1x core"][0]).point

    fig, ax = plt.subplots(figsize=(6.0, 2.4))
    ax.axhline(control, ls="--", color="0.5", lw=0.8,
               label="non-neutral control")
    ax.errorbar(depths, points, yerr=[err_lo, err_hi], fmt="o",
                ms=5, color="#16407a", capsize=3, lw=0.8,
                label="neutral, Wilson 95% CI")
    ax.set_yscale("log")
    ax.set_xlabel("Unitary-core iterations")
    ax.set_ylabel(r"Dump probability $\hat p$")
    ax.set_xticks([0, 1, 2, 3])
    ax.set_xlim(-0.5, 3.5)
    ax.set_ylim(2e-5, 5e-1)
    ax.grid(True, which="major", color="0.9")
    ax.legend(frameon=False, fontsize=8, loc="upper right")
    fig.tight_layout()

    png, pdf = f"{path_stem}.png", f"{path_stem}.pdf"
    fig.savefig(png, dpi=200)
    fig.savefig(pdf)
    plt.close(fig)
    return png


def fig_exp2(path_stem: str = "fig_exp2_three_tier") -> str:
    """Bar chart of the Experiment 2 three-tier neutral-heralding result.

    Dump probability per input on a log scale with 95% Wilson error bars,
    colour-coded by tier (non-neutral control / neutral inputs / pure DC) and
    annotated with the ideal 1/8 dump line.
    """
    tier_color = {
        "control": "#d9822b",   # amber
        "neutral": "#16407a",   # blue
        "dc": "#2e7d32",        # green
    }
    tier_label = {
        "control": "non-neutral control",
        "neutral": "neutral inputs",
        "dc": "pure DC",
    }

    labels, points, lo, hi, colors = [], [], [], [], []
    floor = 1e-4  # log-axis clamp for the ~2e-4 point
    for r in hw.EXP2:
        ci = wilson_ci(r.x, r.n)
        labels.append(r.label)
        points.append(max(ci.point, floor))
        lo.append(max(ci.point, floor) - max(ci.low, floor / 2))
        hi.append(ci.high - max(ci.point, floor))
        colors.append(tier_color.get(r.kind, "0.5"))

    x = np.arange(len(labels))
    fig, ax = plt.subplots(figsize=(6.4, 3.2))
    ax.bar(x, points, color=colors, width=0.62, zorder=2)
    ax.errorbar(x, points, yerr=[lo, hi], fmt="none", ecolor="0.2",
                capsize=3, lw=0.8, zorder=3)
    ax.axhline(0.125, ls="--", color="0.45", lw=0.8, zorder=1,
               label="ideal dump 1/8")

    ax.set_yscale("log")
    ax.set_ylim(floor / 2, 1.5)
    ax.set_ylabel(r"Dump probability $\hat p$")
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=8)
    ax.grid(True, axis="y", which="major", color="0.9", zorder=0)

    handles = [plt.Rectangle((0, 0), 1, 1, color=c) for c in tier_color.values()]
    handles.append(plt.Line2D([0], [0], ls="--", color="0.45"))
    ax.legend(handles, list(tier_label.values()) + ["ideal dump 1/8"],
              frameon=False, fontsize=8, loc="lower left", ncol=2)
    fig.tight_layout()

    png, pdf = f"{path_stem}.png", f"{path_stem}.pdf"
    fig.savefig(png, dpi=200)
    fig.savefig(pdf)
    plt.close(fig)
    return png


if __name__ == "__main__":
    for out in (fig_exp4(), fig_exp2()):
        print(f"wrote {out} (and .pdf)")
