"""Generate constraint annealing visual for docs."""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch
from matplotlib.gridspec import GridSpec

# ---------- palette ----------
FROZEN  = "#2166ac"
WARM    = "#f4a582"
HOT     = "#b2182b"
BG      = "#fafafa"
INIT_C  = "#4393c3"
SYNTH_C = "#7570b3"
ACCUM_C = "#1b9e77"
GRAD_C  = "#d95f02"
GREY    = "#555555"
WARM_TXT = "#c06030"      # darker warm for readable text

def boltzmann(dE, kT):
    return np.exp(-dE / np.clip(kT, 1e-6, None))

def rbox(ax, cx, cy, w, h, text, fc, fs=9.5, tc="white", ec=None):
    """Rounded box with centred label."""
    ax.add_patch(FancyBboxPatch(
        (cx - w/2, cy - h/2), w, h, boxstyle="round,pad=0.12",
        fc=fc, ec=ec or fc, alpha=0.92, lw=1.8, zorder=3))
    ax.text(cx, cy, text, ha="center", va="center", fontsize=fs,
            color=tc, fontweight="bold", zorder=4, linespacing=1.25)

def arrow(ax, a, b, color=GREY, lw=1.8, head=True, ls="solid"):
    ax.add_patch(FancyArrowPatch(
        a, b, arrowstyle="-|>" if head else "-", color=color,
        lw=lw, linestyle=ls, mutation_scale=14, zorder=2))

# ── figure ───────────────────────────────────────────────────────
fig = plt.figure(figsize=(19, 7), facecolor="white")
gs  = GridSpec(1, 2, width_ratios=[0.8, 1.5], wspace=0.25)

# ================================================================
# LEFT — Boltzmann curves
# ================================================================
ax = fig.add_subplot(gs[0])
dE = np.linspace(0, 8, 300)
for kT, c, lb in [(0.3, FROZEN, r"$T\!\approx\!0$  (frozen)"),
                   (1.5, WARM,   r"$T=1$  (warm)"),
                   (5.0, HOT,    r"$T\!\gg\!1$  (hot)")]:
    ax.plot(dE, boltzmann(dE, kT), c=c, lw=2.5, label=lb)
    ax.fill_between(dE, boltzmann(dE, kT), alpha=.08, color=c)
ax.set(xlabel=r"Deviation from plan  $\Delta E$",
       ylabel=r"Acceptance  $P = e^{-\Delta E/kT}$",
       xlim=(0, 8), ylim=(0, 1.05))
ax.xaxis.label.set_size(12); ax.yaxis.label.set_size(12)
ax.set_title("Boltzmann Analogy", fontsize=13, fontweight="bold", pad=10)
ax.legend(fontsize=10, loc="upper right", framealpha=.9)
ax.set_facecolor(BG); ax.grid(True, alpha=.3, ls="--")
ax.annotate("Only parameter\nfixes allowed",
    xy=(1, boltzmann(1, .3)), xytext=(3, .55), fontsize=9,
    color=FROZEN, ha="center",
    arrowprops=dict(arrowstyle="->", color=FROZEN, lw=1.2))
ax.annotate("Model changes\nif justified",
    xy=(3, boltzmann(3, 1.5)), xytext=(5.5, .65), fontsize=9,
    color=WARM_TXT, ha="center",
    arrowprops=dict(arrowstyle="->", color=WARM_TXT, lw=1.2))
ax.annotate("Full freedom\n(justify from data)",
    xy=(5, boltzmann(5, 5)), xytext=(6.5, .5), fontsize=9,
    color=HOT, ha="center",
    arrowprops=dict(arrowstyle="->", color=HOT, lw=1.2))

# ================================================================
# RIGHT — Verification & Learning Loop
# ================================================================
ax2 = fig.add_subplot(gs[1])
ax2.set_facecolor(BG); ax2.axis("off")
ax2.set_title("Verification & Learning Loop  (max_iter=5)",
              fontsize=13, fontweight="bold", pad=12)

# ── bar geometry ─────────────────────────────────────────────────
cx_bars = 2.5
bar_w   = 1.7
bar_h   = 0.42
gap     = 0.16
top_y   = 4.5

def yb(i):
    return top_y - i * (bar_h + gap)

levels = [0, 0, 1, 1, 2]
cmap   = [FROZEN, WARM, HOT]
lnames = ["Frozen", "Warm", "Hot"]

# ── Agent Init ───────────────────────────────────────────────────
init_cy = top_y + 0.80
rbox(ax2, cx_bars, init_cy, 1.35, 0.50, "Agent Init", INIT_C, fs=11)
arrow(ax2, (cx_bars, init_cy - 0.27), (cx_bars, yb(0) + bar_h/2 + 0.08))

# ── verification group background ────────────────────────────────
vpad = 0.25
ax2.add_patch(FancyBboxPatch(
    (cx_bars - bar_w/2 - vpad, yb(4) - bar_h/2 - vpad),
    bar_w + 2*vpad, (yb(0) + bar_h/2 + vpad) - (yb(4) - bar_h/2 - vpad),
    boxstyle="round,pad=0.18", fc="#e2e2e2", ec="#c0c0c0",
    alpha=0.55, lw=1.2, zorder=0))

# ── temperature bars ─────────────────────────────────────────────
for i, lvl in enumerate(levels):
    cy = yb(i)
    ax2.barh(cy, bar_w, left=cx_bars - bar_w/2, height=bar_h,
             color=cmap[lvl], alpha=.87, edgecolor="white", lw=2, zorder=1)
    ax2.text(cx_bars - bar_w/2 - 0.15, cy, str(i),
             ha="center", va="center", fontsize=10.5, fontweight="bold",
             color="#333")
    ax2.text(cx_bars, cy, f"T={lvl}  ({lnames[lvl]})",
             ha="center", va="center", fontsize=10,
             color="white" if lvl != 1 else "#333", fontweight="bold",
             zorder=2)

# ── descriptions: one per temperature group, left-aligned ────────
# Centred on the vertical midpoint of each bar group
group_descs = [
    ([0, 1], "Fix parameters, bounds,\nbaseline, convergence",   FROZEN),
    ([2, 3], "Prefer smallest change;\ndeviate if justified",    WARM_TXT),
    ([4],    "Full freedom;\njustify from data",                 HOT),
]
desc_x = cx_bars - bar_w/2 - 0.28
for rows, txt, col in group_descs:
    mid_y = sum(yb(r) for r in rows) / len(rows)
    ax2.text(desc_x, mid_y, txt, ha="right", va="center",
             fontsize=8.5, color=col, linespacing=1.35, fontstyle="italic")

# ── "temperature rises" arrow ────────────────────────────────────
arr_x = 0.15
ax2.annotate("", xy=(arr_x, yb(4) - 0.12), xytext=(arr_x, yb(0) + 0.12),
    arrowprops=dict(arrowstyle="-|>", color=GREY, lw=2))
ax2.text(arr_x - 0.13, (yb(0) + yb(4)) / 2,
         "Temperature rises", ha="center", va="center",
         fontsize=9, color=GREY, rotation=90, fontstyle="italic")

# ── Knowledge Synthesis ──────────────────────────────────────────
synth_cy = yb(4) - 0.90
sw, sh   = 1.55, 0.50
rbox(ax2, cx_bars, synth_cy, sw, sh, "Knowledge Synthesis", SYNTH_C, fs=10)
arrow(ax2, (cx_bars, yb(4) - bar_h/2 - 0.08), (cx_bars, synth_cy + sh/2 + 0.05))

# ── Accumulated Knowledge (right side, level with bars 3-4) ──────
accum_cx = 4.45
accum_cy = (yb(3) + yb(4)) / 2
aw, ah   = 1.30, 0.58
rbox(ax2, accum_cx, accum_cy, aw, ah, "Accumulated\nKnowledge", ACCUM_C, fs=10)

# ── Graduation to Skills (branches right from Accumulated) ───────
grad_cx = accum_cx + 1.70
grad_cy = accum_cy
gw, gh  = 1.20, 0.58
rbox(ax2, grad_cx, grad_cy, gw, gh, "Graduation\nto Skills", GRAD_C, fs=10)

# ── connectors (teal loop) ──────────────────────────────────────

# Knowledge Synthesis → right → up → Accumulated Knowledge  (short L)
corner_br = (accum_cx, synth_cy)
arrow(ax2, (cx_bars + sw/2 + 0.05, synth_cy), corner_br,
      color=ACCUM_C, lw=2.2, head=False)
arrow(ax2, corner_br, (accum_cx, accum_cy - ah/2 - 0.05),
      color=ACCUM_C, lw=2.2)

# Accumulated Knowledge → up → left → Agent Init  (tall L)
corner_tr = (accum_cx, init_cy)
arrow(ax2, (accum_cx, accum_cy + ah/2 + 0.05), corner_tr,
      color=ACCUM_C, lw=2.2, head=False)
arrow(ax2, corner_tr, (cx_bars + 0.70, init_cy),
      color=ACCUM_C, lw=2.2)

# Accumulated Knowledge → Graduation to Skills  (dashed branch)
arrow(ax2, (accum_cx + aw/2 + 0.05, accum_cy),
           (grad_cx - gw/2 - 0.05, grad_cy),
      color=GRAD_C, lw=1.8, ls=(0, (4, 3)))

ax2.text(grad_cx, grad_cy - gh/2 - 0.10,
         "(when sufficient evidence)", ha="center", va="top",
         fontsize=8.5, color=GRAD_C, fontstyle="italic")

# ── axis limits ──────────────────────────────────────────────────
ax2.set_xlim(-0.15, grad_cx + gw/2 + 0.15)
ax2.set_ylim(synth_cy - 0.55, init_cy + 0.55)

# ── main title ───────────────────────────────────────────────────
fig.suptitle(
    r"Constraint Annealing:  $P \propto e^{-\Delta E/kT}$"
    "  applied to plan conformance",
    fontsize=15, fontweight="bold", y=0.98)

plt.savefig("docs/constraint_annealing.png", dpi=180,
            bbox_inches="tight", facecolor="white")
plt.close()
print("Saved docs/constraint_annealing.png")
