"""T140 original figures from finite catalogue outputs; no hardware ranking."""
import os
from pathlib import Path
os.environ.setdefault("MPLCONFIGDIR", str(Path(__file__).resolve().parents[2]/"evidence/quantum-memory-series/work/mpl-cache"))
import csv
import json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import TwoSlopeNorm
import matched_protection_budgets as m

plt.rcParams.update({"font.family": ["DejaVu Sans", "WenQuanYi Zen Hei"], "font.size": 11,
                     "axes.unicode_minus": False, "axes.spines.top": False, "axes.spines.right": False,
                     "figure.facecolor": "#fafcfb", "axes.facecolor": "#fafcfb", "savefig.facecolor": "#fafcfb"})
teal, orange, blue, grey = "#237c73", "#ce774b", "#487996", "#adc9c0"
assert json.loads((m.OUT/"audit.json").read_text())["passed"]
r = json.loads((m.OUT/"results.json").read_text())
read = lambda name: list(csv.DictReader((m.OUT/name).open()))
fronts = read("matched-budget-envelopes.csv")

def point(W, kz, Bg, Br):
    return next(a for a in fronts if all(float(a[k]) == v for k, v in [
        ("deadline", W), ("kz", kz), ("coupling_cap", Bg), ("aux_reset_cap", Br)]))

fig, axs = plt.subplots(1, 2, figsize=(11.8, 5.2))
best = np.zeros((7, 7)); diff = np.full((7, 7), np.nan)
for y, Br in enumerate(m.R_BUDGETS):
    for x, Bg in enumerate(m.G_BUDGETS):
        a = point(2., .02, Bg, Br)
        best[y, x] = float(a["best_Fe"])
        if a["continuous_Fe"] and a["switched_Fe"]:
            diff[y, x] = float(a["continuous_Fe"])-float(a["switched_Fe"])
im = axs[0].imshow(best, origin="lower", aspect="auto", cmap="YlGnBu", vmin=.68, vmax=max(.84, best.max()))
span = float(np.nanmax(abs(diff)))
im2 = axs[1].imshow(diff, origin="lower", aspect="auto", cmap="RdBu", norm=TwoSlopeNorm(vcenter=0, vmin=-span, vmax=span))
for ax in axs:
    ax.set(xticks=range(7), xticklabels=[f"{x:g}" for x in m.G_BUDGETS],
           yticks=range(7), yticklabels=[f"{x:g}" for x in m.R_BUDGETS],
           xlabel="耦合强度平方积分上限", ylabel="辅助复位强度积分上限（含准备）")
for y in range(7):
    for x in range(7):
        axs[0].text(x, y, f"{best[y,x]:.3f}", ha="center", va="center", fontsize=8,
                    color="white" if best[y,x]>.79 else "#253b34")
        if np.isfinite(diff[y,x]):
            label = "0.000" if abs(diff[y,x]) < .0005 else f"{diff[y,x]:+.3f}"
            axs[1].text(x, y, label, ha="center", va="center", fontsize=8,
                        color="white" if abs(diff[y,x]) > .55*span else "#253b34")
axs[0].set_title("全部已执行可行方案中的最高 Fe")
axs[1].set_title("连续 Fe − 分段 Fe；空格为不具备两类候选")
fig.colorbar(im, ax=axs[0], fraction=.046, pad=.03); fig.colorbar(im2, ax=axs[1], fraction=.046, pad=.03)
fig.suptitle("同物理输入/输出、W=2、X率0.2/Z率0.02；离散预算格点，不是能耗评分", fontsize=12)
fig.tight_layout(); fig.savefig(m.OUT/"separate-resource-caps.png", dpi=165); plt.close(fig)

fig, axs = plt.subplots(1, 2, figsize=(11, 4.8))
styles = [("bare_Fe", "#58675f", "裸比特"), ("encoded_no_repair_Fe", grey, "编码，无持续修复"),
          ("continuous_Fe", teal, "连续修复"), ("switched_Fe", orange, "分段修复")]
for key, col, label in styles:
    rows = [point(W, .02, 80., 40.) for W in m.DEADLINES]
    axs[0].plot(m.DEADLINES, [float(a[key]) for a in rows], "o-", color=col, label=label)
    rows = [point(2., kz, 80., 40.) for kz in m.Z_RATES]
    axs[1].plot(m.Z_RATES, [float(a[key]) for a in rows], "o-", color=col, label=label)
for ax in axs:
    ax.set(ylabel="无条件物理输入输出 Fe", ylim=(.38, 1.))
    ax.legend(fontsize=8); ax.grid(alpha=.15)
axs[0].set(xlabel="共同截止时间 W", title="固定两项控制上限 80、40；Z率0.02")
axs[1].set(xlabel="每数据比特的未修复 Z 噪声率", title="同 W=2；保护收益会随噪声方向改变")
fig.suptitle("有限候选包络；关键预算点另作局部细化，均不代表全控制最优", fontsize=12)
fig.tight_layout(); fig.savefig(m.OUT/"deadline-and-noise-boundaries.png", dpi=165); plt.close(fig)

pareto = read("quality-resource-frontiers.csv")
fig, axs = plt.subplots(1, 2, figsize=(11, 4.7))
for ax, kz, target in zip(axs, [0., .02], [.9, .8]):
    for kind, col, label in [("continuous", teal, "连续"), ("switched", orange, "分段")]:
        rows = [a for a in pareto if float(a["deadline"]) == 2. and float(a["kz"]) == kz
                and float(a["target_Fe"]) == target and a["kind"] == kind]
        ax.plot([float(a["coupling_integral"]) for a in rows],
                [float(a["aux_reset_integral"]) for a in rows], "o-", color=col, markersize=3, label=label)
    ax.set(xlabel="实际使用的耦合强度平方积分", ylabel="实际使用的辅助复位强度积分",
           title=f"W=2，Z率{kz:g}，要求 Fe≥{target:g}")
    ax.legend(fontsize=9); ax.grid(alpha=.15)
fig.suptitle("达到同一质量的两项资源前沿；连线只引导视线，不补算中间方案", fontsize=12)
fig.tight_layout(); fig.savefig(m.OUT/"quality-resource-frontiers.png", dpi=165); plt.close(fig)

key = "W2_Z0.02_G80_R40"
ex = r["examples"][key]["protocols"]
fig, axs = plt.subplots(2, 1, figsize=(11, 5.5), gridspec_kw={"height_ratios": [1, 1.45]})
names = ["encoded_no_repair", "continuous", "switched"]
labels = ["编码，无持续修复", "连续修复", "分段修复"]
for i, name in enumerate(names):
    a = ex[name]["metrics"]; left = 0.
    for width, col in [(a["prep"], grey), (.04, blue), (a["storage_core_time"], teal), (.08, orange)]:
        axs[0].barh(i, width, left=left, color=col, height=.6); left += width
axs[0].set(yticks=range(3), yticklabels=labels, xlim=(0, 2), xlabel="准备 → 编码 → 有限控制存储 → 解码；共同截止 W=2")
axs[1].axis("off")
for i, name in enumerate(names):
    a = ex[name]["metrics"]
    text = (f"{labels[i]}：Fe {a['Fe']:.4f}；耦合积分 {a['coupling_budget_used']:.2f}；"
            f"复位积分 {a['aux_reset_budget_used']:.2f}\n"
            f"存储时平均实际复位事件 {a['core_aux_reset_events']:.3f}；"
            f"末辅助激发概率 {a['final_aux_excitation']:.3f}")
    axs[1].text(.02, .94-i*.31, text, va="top", fontsize=11, color="#304b45")
fig.suptitle("同任务的时钟账与资源账分别记录；辅助末状态不自动等于可复用", fontsize=12)
fig.tight_layout(); fig.savefig(m.OUT/"physical-service-ledger.png", dpi=165); plt.close(fig)

fig, ax = plt.subplots(figsize=(12, 6.2)); ax.axis("off")
ax.text(.07, .8, "先对齐任务，再比较保护收益", fontsize=27, color=teal)
ax.text(.08, .54, "同一输入与输出 · 同一截止时间", fontsize=23, color=teal)
ax.text(.08, .31, "耦合预算与复位预算，分别记账", fontsize=23, color=orange)
ax.text(.08, .08, "有限未知态存储模型；完整准备、编码、恢复与解码", fontsize=13, color="#647c73")
fig.savefig(m.OUT/"cover.png", dpi=140, bbox_inches="tight"); plt.close(fig)
print("saved 4 body figures and cover")