"""Generate figure: constraint annealing in the continuous learning loop.

Follows the ASCII layout exactly:
  - Row 1: Analysis (left) → Annealing (center) → Result (right)
  - Row 2: ↑ return line (left), Knowledge Base (center) ← Knowledge Synthesis (right)
  - Row 3: ↑ return line (left), Skill Graduation (center) → Domain Skill (right)
  - Return: Domain Skill → down → left → up → Analysis
"""

import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch

# ── palette ──
ANALYSIS = "#4393c3"
FROZEN = "#2166ac"
WARM = "#f4a582"
HOT = "#b2182b"
ANNEAL_BG = "#efefef"
KNOWLEDGE = "#7b3294"
SKILL = "#d6604d"
ARROW = "#444"
GRAY = "#777"


def box(ax, cx, cy, w, h, text, color, fontsize=10, fc="white"):
    """Rounded box centered at (cx, cy)."""
    x0, y0 = cx - w / 2, cy - h / 2
    b = FancyBboxPatch(
        (x0, y0), w, h, boxstyle="round,pad=0.12",
        facecolor=color, edgecolor="white", alpha=0.9, lw=2,
    )
    ax.add_patch(b)
    ax.text(cx, cy, text, ha="center", va="center",
            fontsize=fontsize, color=fc, fontweight="bold", linespacing=1.3)
    return x0, y0, w, h


def arr(ax, x0, y0, x1, y1, color=ARROW, lw=1.8, hw=0.12, hl=0.15):
    """Arrow from (x0,y0) to (x1,y1)."""
    ax.annotate("", xy=(x1, y1), xytext=(x0, y0),
                arrowprops=dict(arrowstyle=f"-|>,head_width={hw},head_length={hl}",
                                color=color, lw=lw))


def label(ax, x, y, text, color=GRAY, fontsize=8.5, ha="center", va="center"):
    ax.text(x, y, text, fontsize=fontsize, color=color, fontstyle="italic",
            ha=ha, va=va, linespacing=1.3)


# ── layout grid (center coords) ──
LEFT_X = 2.0    # Analysis column
MID_X = 7.0     # Annealing / Knowledge Base / Skill Graduation
RIGHT_X = 12.0  # Result / Knowledge Synthesis / Domain Skill

ROW1_Y = 8.0    # top row
ROW2_Y = 5.0    # middle row
ROW3_Y = 2.0    # bottom row

BW, BH = 3.0, 1.3  # box width, height

fig, ax = plt.subplots(figsize=(16, 11), facecolor="white")
ax.set_xlim(-1, 16)
ax.set_ylim(-0.5, 10.5)
ax.set_aspect("equal")
ax.axis("off")

# ══════════════════════════════════════════════════════════
#  Title + banner
# ══════════════════════════════════════════════════════════
ax.text(7.5, 10.2,
        "Constraint Annealing in the Continuous Learning Loop",
        ha="center", va="center", fontsize=15, fontweight="bold")

banner = FancyBboxPatch(
    (0.5, 9.35), 14, 0.55, boxstyle="round,pad=0.1",
    facecolor="#f7f7f7", edgecolor="#ccc", lw=1,
)
ax.add_patch(banner)
ax.text(7.5, 9.62,
        "Better skills  \u2192  better plans  \u2192  less annealing  "
        "\u2192  higher quality  \u2192  better knowledge  \u2192  better skills",
        ha="center", va="center", fontsize=9, color="#555", fontstyle="italic")

# ══════════════════════════════════════════════════════════
#  ROW 1: Analysis → Annealing → Result
# ══════════════════════════════════════════════════════════
box(ax, LEFT_X, ROW1_Y, BW, BH,
    "Analysis\n(Curve Fitting Agent)", ANALYSIS, fontsize=11)

# Annealing box (special: light bg with mini bars)
ann_w, ann_h = 3.4, 1.5
ann_x0 = MID_X - ann_w / 2
ann_y0 = ROW1_Y - ann_h / 2
ax.add_patch(FancyBboxPatch(
    (ann_x0, ann_y0), ann_w, ann_h, boxstyle="round,pad=0.1",
    facecolor=ANNEAL_BG, edgecolor="#bbb", lw=1.2,
))
ax.text(MID_X, ROW1_Y + ann_h / 2 - 0.18,
        "Constraint Annealing", ha="center", va="top",
        fontsize=10, fontweight="bold", color="#333")
ax.text(MID_X, ROW1_Y + ann_h / 2 - 0.45,
        r"$P \propto e^{-\Delta E\,/\,kT}$", ha="center", va="top",
        fontsize=9.5, color="#555")

bars = [("T=0  Frozen", FROZEN, "white"),
        ("T=1  Warm", WARM, "#333"),
        ("T=2  Hot", HOT, "white")]
for j, (lbl, clr, fc) in enumerate(bars):
    by = ROW1_Y + 0.08 - j * 0.32
    ax.barh(by, 2.6, left=ann_x0 + 0.4, height=0.25,
            color=clr, alpha=0.85, edgecolor="white", lw=1)
    ax.text(ann_x0 + 0.4 + 1.3, by, lbl, ha="center", va="center",
            fontsize=7.5, color=fc, fontweight="bold")

box(ax, RIGHT_X, ROW1_Y, BW, BH,
    "Result +\nquality_history", ANALYSIS, fontsize=10, fc="white")

# quality_history bullets
qh = ["- verif iterations", "- alt models", "- script errors", "- judge reasoning"]
for k, item in enumerate(qh):
    ax.text(RIGHT_X + BW / 2 + 0.2, ROW1_Y + 0.45 - k * 0.3, item,
            fontsize=7.5, color="#555", va="center", ha="left")

# Arrows: Analysis → Annealing → Result
arr(ax, LEFT_X + BW / 2, ROW1_Y, ann_x0, ROW1_Y)
arr(ax, ann_x0 + ann_w, ROW1_Y, RIGHT_X - BW / 2, ROW1_Y)

# ══════════════════════════════════════════════════════════
#  ROW 2: Knowledge Base (center) ← Knowledge Synthesis (right)
# ══════════════════════════════════════════════════════════
box(ax, MID_X, ROW2_Y, BW, BH,
    "Knowledge Base\n(accumulated)", KNOWLEDGE, fontsize=10)

box(ax, RIGHT_X, ROW2_Y, BW, BH,
    "Knowledge Synthesis\n(failure / method)", KNOWLEDGE, fontsize=10)

# Arrow ↓ Result → Knowledge Synthesis
arr(ax, RIGHT_X, ROW1_Y - BH / 2, RIGHT_X, ROW2_Y + BH / 2)
label(ax, RIGHT_X + 1.6, (ROW1_Y + ROW2_Y) / 2,
      "quality_history\nfeeds synthesis", color=KNOWLEDGE)

# Arrow ← Knowledge Synthesis → Knowledge Base
arr(ax, RIGHT_X - BW / 2, ROW2_Y, MID_X + BW / 2, ROW2_Y)

# Arrow ← Knowledge Base → return line (left)
arr(ax, MID_X - BW / 2, ROW2_Y, LEFT_X, ROW2_Y, color=KNOWLEDGE, lw=2)

# Return line up: left vertical arrow to Analysis
arr(ax, LEFT_X, ROW2_Y, LEFT_X, ROW1_Y - BH / 2, color=KNOWLEDGE, lw=2)
label(ax, LEFT_X - 1.2, (ROW1_Y + ROW2_Y) / 2,
      "prior knowledge\ninjected", color=KNOWLEDGE)

# ══════════════════════════════════════════════════════════
#  ROW 3: Skill Graduation (center) → Domain Skill (right)
# ══════════════════════════════════════════════════════════
box(ax, MID_X, ROW3_Y, BW, BH,
    "Skill Graduation\n(LLM distills)", SKILL, fontsize=10)

box(ax, RIGHT_X, ROW3_Y, BW, BH,
    "Domain Skill (.md)\nplan / analysis /\ninterp / validation", SKILL, fontsize=9)

# Arrow ↓ Knowledge Base → Skill Graduation
arr(ax, MID_X, ROW2_Y - BH / 2, MID_X, ROW3_Y + BH / 2, color=KNOWLEDGE)
label(ax, MID_X + 1.8, (ROW2_Y + ROW3_Y) / 2,
      "enough evidence\naccumulated", color=KNOWLEDGE)

# Arrow → Skill Graduation → Domain Skill
arr(ax, MID_X + BW / 2, ROW3_Y, RIGHT_X - BW / 2, ROW3_Y)

# ══════════════════════════════════════════════════════════
#  Return loop: Domain Skill → down → left → up → Analysis
# ══════════════════════════════════════════════════════════
RETURN_Y = 0.5  # bottom of the return path

# Down from Domain Skill
ax.plot([RIGHT_X, RIGHT_X], [ROW3_Y - BH / 2, RETURN_Y],
        color=SKILL, lw=2.5, solid_capstyle="butt")
# Left along bottom
ax.plot([RIGHT_X, LEFT_X], [RETURN_Y, RETURN_Y],
        color=SKILL, lw=2.5, solid_capstyle="butt")
# Up to Analysis (with arrowhead)
arr(ax, LEFT_X, RETURN_Y, LEFT_X, ROW2_Y - 0.05, color=SKILL, lw=2.5)

# Label on return path
label(ax, (LEFT_X + RIGHT_X) / 2, RETURN_Y - 0.35,
      "MANDATORY skill rules injected at every stage  (reduces need for annealing)",
      color=SKILL, fontsize=9)

# update_skill note
label(ax, (MID_X + RIGHT_X) / 2, ROW3_Y - BH / 2 - 0.3,
      "update_skill() when new evidence arrives", color=GRAY, fontsize=8)

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