"""T139 original figures from saved finite calculations, no experimental images."""
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
import error_corrected_sensing 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"
r = json.loads((m.OUT/"results.json").read_text())
assert json.loads((m.OUT/"audit.json").read_text())["passed"]
read = lambda name: list(csv.DictReader((m.OUT/name).open()))

fig, axs = plt.subplots(1, 2, figsize=(10.8, 4.5))
for ax, row, label in zip(axs, r["opening"], ["信号 X：逻辑空间保留响应", "信号 Z：逻辑响应消失，记录仍有信息"]):
    values = [row[k] for k in ["CQFI", "record_ignored_QFI", "record_FI"]]
    ax.bar(range(3), values, color=[teal, blue, orange])
    for i, v in enumerate(values): ax.text(i, v+max(values)*.025, f"{v:.4f}", ha="center")
    ax.set(xticks=range(3), xticklabels=["记录＋量子输出", "忽略记录", "仅记录"], ylabel="局部 Fisher 信息", title=label, ylim=(0, max(values)*1.2))
fig.suptitle("同一编码与 Z 噪声，改变待测信号；全部 32 个分支均保留", fontsize=13)
fig.tight_layout(); fig.savefig(m.OUT/"signal-and-records.png", dpi=165); plt.close(fig)

markov, bare = read("markov-control-scan.csv"), read("bare-initial-time-scan.csv")
fig, axs = plt.subplots(1, 2, figsize=(11, 4.7))
for ax, wall in zip(axs, [1.1, 4.1]):
    for latency, col in zip([0., .005, .02, .05], [teal, blue, orange, grey]):
        rows = [d for d in markov if d["signal"] == "X" and abs(float(d["wall_time"])-wall)<1e-8 and abs(float(d["latency_per_cycle"])-latency)<1e-8]
        ax.plot([int(d["cycles"]) for d in rows], [float(d["actual_local_readout_FI"]) for d in rows], "o-", color=col, label=f"每轮等待 {latency:g}")
    b = next(d for d in bare if d["signal"] == "X" and abs(float(d["sensing_time"])-(wall-.1))<1e-8)
    ax.axhline(float(b["actual_local_readout_FI"]), color="#4b5652", ls="--", label="优化裸探针初态")
    ax.set(xscale="log", xlabel="纠错次数（信号开启时间随次数缩短）", ylabel="实际局部读出 FI", title=f"同一总时长 {wall:g}；准备/读出 0.1")
    ax.legend(fontsize=8); ax.grid(alpha=.15)
fig.suptitle("有限控制时间使纠错更密未必更好；末读出翻转率 0.02", fontsize=13)
fig.tight_layout(); fig.savefig(m.OUT/"finite-control-budget.png", dpi=165); plt.close(fig)

hmm = read("hidden-environment-scan.csv")
fig, axs = plt.subplots(1, 2, figsize=(11, 4.7))
rows = [d for d in hmm if d["initial_environment"] == "[0.5, 0.5]" and abs(float(d["sensing_time"])-2.)<1e-8]
for key, col, label in [("CQFI", teal, "全部可见记录＋量子输出"), ("record_ignored_QFI", orange, "忽略记录"), ("environment_exposed_QFI_bound", blue, "额外暴露环境的上界")]:
    axs[0].plot([float(d["nu"]) for d in rows], [float(d[key]) for d in rows], "o-", color=col, label=label)
axs[0].set(xscale="symlog", xlabel="环境切换速率 ν", ylabel="局部 QFI", title="环境边际始终 1/2，时间记忆不同")
axs[0].legend(fontsize=9); axs[0].grid(alpha=.15)
for p0, col in [(1., teal), (.5, blue), (0., orange)]:
    rows = [d for d in hmm if d["initial_environment"] == str([p0, 1-p0]) and abs(float(d["sensing_time"])-2.)<1e-8]
    axs[1].plot([float(d["nu"]) for d in rows], [float(d["actual_local_readout_FI"]) for d in rows], "o-", color=col, label=f"初始好环境概率 {p0:g}")
axs[1].set(xscale="symlog", xlabel="环境切换速率 ν", ylabel="实际局部读出 FI", title="未访问的环境扇区也会改变一般条件")
axs[1].legend(fontsize=9); axs[1].grid(alpha=.15)
fig.suptitle("本库两态隐藏环境模型；8 次理想检查，等待 0.02，环境从未中途重置", fontsize=12)
fig.tight_layout(); fig.savefig(m.OUT/"hidden-environment-memory.png", dpi=165); plt.close(fig)

fock, optima = read("fock-repair-scan.csv"), read("fock-time-optima.csv")
fig, axs = plt.subplots(1, 2, figsize=(11, 4.7))
for repair, col in [(0., orange), (1., grey), (10., blue), (100., teal)]:
    rows = [d for d in fock if float(d["repair_over_kappa"]) == repair]
    axs[0].plot([float(d["t"]) for d in rows], [float(d["FI_rate"]) for d in rows], "o-", color=col, label=f"修复 Γ/κ = {repair:g}")
t = np.linspace(.01, 10., 200)
axs[0].plot(t, 9*t*t*np.exp(-t)/(t+.1), "--", color="#465b51", label="无限快修复极限，仍衰减")
axs[0].set(xlabel="询问时间（κ⁻¹）", ylabel="FI / 总时长", title="同一 Fock1,4 初态与完整输出")
axs[0].legend(fontsize=8); axs[0].grid(alpha=.15)
rows = [d for d in fock if float(d["repair_over_kappa"]) == 0.]
axs[1].plot([float(d["t"]) for d in rows], [float(d["QFI"]) for d in rows], "o-", color=teal, label="全部损失空间的 QFI / 可达读出 FI")
axs[1].plot([float(d["t"]) for d in rows], [float(d["code_only_Ramsey_FI"]) for d in rows], "o-", color=orange, label="只在原码空间作 Ramsey 读出")
axs[1].set(xlabel="未保护询问时间（κ⁻¹）", ylabel="局部 FI", title="把损失空间丢掉，会压低裸基线")
axs[1].legend(fontsize=8); axs[1].grid(alpha=.15)
fig.suptitle("自建连续五能级共享泵模型；不拟合实际腔实验，泵入光子另计", fontsize=12)
fig.tight_layout(); fig.savefig(m.OUT/"approximate-fock-sensing.png", dpi=165); plt.close(fig)

fig, ax = plt.subplots(figsize=(12, 6.2)); ax.axis("off")
ax.text(.07, .79, "纠错要留下信号，也要留下可用读出", fontsize=25, color=teal)
ax.text(.08, .52, "噪声可修复  →  信号仍能积累", fontsize=23, color=teal)
ax.text(.08, .29, "逻辑响应消失  ≠  全部记录没有信息", fontsize=22, color=orange)
ax.text(.08, .07, "有限控制、隐藏环境与近似 Fock 编码的可回查例子", 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")