docs/make_figures.py

"""
Ve tat ca so do cho chuong Methodology -> file PNG trong docs/figures/.
Phong cach: den trang, toi gian, it chu. Frozen = o xam nhat, Trainable = o trang.
Rieng fig01 (so do tong quan) ve "dep" kieu paper, co mau + icon.

CHAY: bam Run (hoac: python make_figures.py). Anh xuat ra docs/figures/.
"""

import math
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle, FancyBboxPatch, Ellipse, Circle, Polygon
from pathlib import Path

plt.rcParams.update({"font.size": 9, "font.family": "DejaVu Sans"})

OUT = Path(__file__).resolve().parent / "figures"
OUT.mkdir(exist_ok=True)

GRAY = "0.88"   # frozen (cac hinh don gian)

# ============================================================================
#  CONFIG cho fig01: anh chest X-ray that o o input.
#  - Doi sang anh cua ban neu muon (duong dan tuyet doi).
#  - De None -> ve icon X-quang cach dieu thay cho anh that.
# ============================================================================
CXR_IMAGE = r"D:\USTH\KLTN\cxr-vlm-data\IU-Xray\images\CXR1000_IM-0003-1001.png"


# ---- helpers chung ---------------------------------------------------------

def canvas(W, H):
    fig, ax = plt.subplots(figsize=(W / 10, H / 10))
    ax.set_xlim(0, W)
    ax.set_ylim(0, H)
    ax.set_aspect("equal")
    ax.axis("off")
    return fig, ax


def box(ax, cx, cy, w, h, text, frozen=False, dashed=False):
    fc = GRAY if frozen else "white"
    ax.add_patch(Rectangle((cx - w / 2, cy - h / 2), w, h, facecolor=fc,
                           edgecolor="black", lw=1.2,
                           linestyle="--" if dashed else "-"))
    ax.text(cx, cy, text, ha="center", va="center", fontsize=8.3)
    return cx, cy, w, h


def arr(ax, x1, y1, x2, y2, text=None, dashed=False, color="black", lw=1.1,
        fontsize=7.3):
    ax.annotate("", xy=(x2, y2), xytext=(x1, y1),
                arrowprops=dict(arrowstyle="-|>", color=color, lw=lw,
                                linestyle="--" if dashed else "-",
                                shrinkA=0, shrinkB=0))
    if text:
        ax.text((x1 + x2) / 2, (y1 + y2) / 2 + 0.8, text, ha="center",
                va="bottom", fontsize=fontsize, color=color)


def hrow(ax, cy, specs, h=12):
    """specs: list of (cx, w, text, frozen). Noi cac o bang mui ten ngang."""
    drawn = []
    for cx, w, text, frozen in specs:
        drawn.append(box(ax, cx, cy, w, h, text, frozen=frozen))
    for i in range(len(drawn) - 1):
        cx, cy0, w, _ = drawn[i]
        ncx, _, nw, _ = drawn[i + 1]
        arr(ax, cx + w / 2, cy0, ncx - nw / 2, cy0)
    return drawn


def legend(ax, x, y):
    ax.add_patch(Rectangle((x, y), 5, 3, facecolor=GRAY, edgecolor="black", lw=1))
    ax.text(x + 6, y + 1.5, "Frozen", va="center", fontsize=7.5)
    ax.add_patch(Rectangle((x + 22, y), 5, 3, facecolor="white", edgecolor="black", lw=1))
    ax.text(x + 28, y + 1.5, "Trainable", va="center", fontsize=7.5)


def save(fig, name):
    fig.savefig(OUT / name, dpi=160, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print("[+]", name)


# ---- neuron-style helpers (vong tron = node mang neuron) -------------------

def ncol(ax, x, cy, n, r=1.0, vgap=2.9, fill="white", edge="#2E86C1", z=6):
    """Mot cot 'neuron' (n vong tron) tai hoanh do x, tam doc cy."""
    ys = [cy + (i - (n - 1) / 2) * vgap for i in range(n)]
    for y in ys:
        ax.add_patch(Circle((x, y), r, facecolor=fill, edgecolor=edge, lw=1.0, zorder=z))
    return ys


def nconnect(ax, x1, ys1, x2, ys2, color="#CFCFCF", lw=0.35):
    """Noi day giua 2 cot neuron (fully-connected)."""
    for y1 in ys1:
        for y2 in ys2:
            ax.plot([x1, x2], [y1, y2], color=color, lw=lw, zorder=3)


def rrect(ax, x0, y0, x1, y1, fill="white", edge="#566573", lw=1.4, z=1, dashed=False):
    ax.add_patch(FancyBboxPatch((x0, y0), x1 - x0, y1 - y0,
                 boxstyle="round,pad=0,rounding_size=1.6", facecolor=fill,
                 edgecolor=edge, lw=lw, linestyle="--" if dashed else "-", zorder=z))


# ============================================================================
#  fig01 — so do tong quan kieu paper (mau + icon)
#  CHINH O DAY: doi mau o bang COLORS, doi anh o CXR_IMAGE,
#  doi vi tri/chu o cac lenh trong than ham.
# ============================================================================

def fig01_pipeline():
    FZ_FILL, FZ_EDGE = "#D6EAF8", "#2E86C1"    # frozen  (xanh)
    TR_FILL, TR_EDGE = "#FDEBD0", "#E67E22"    # trainable (cam)
    OUT_FILL, OUT_EDGE = "#E8F8F5", "#16A085"  # output (xanh la)
    ARROW = "#5D6D7E"
    POS, NEG, UNC = "#27AE60", "#C0392B", "#D4A017"  # chip PNU

    W, H = 152, 66
    fig, ax = canvas(W, H)

    # ---- icon helpers ----
    def snow(cx, cy, r, color=FZ_EDGE):
        for k in range(6):
            a = math.radians(k * 60)
            dx, dy = r * math.cos(a), r * math.sin(a)
            ax.plot([cx, cx + dx], [cy, cy + dy], color=color, lw=1.0, zorder=8)
            bx, by = cx + dx * 0.6, cy + dy * 0.6
            for s in (1, -1):
                a2 = a + s * math.radians(38)
                ax.plot([bx, bx + r * 0.32 * math.cos(a2)],
                        [by, by + r * 0.32 * math.sin(a2)], color=color, lw=1.0, zorder=8)

    def flame(cx, cy, s, color=TR_EDGE, inner="#F5B041"):
        body = [(cx, cy + s), (cx + 0.55 * s, cy + 0.1 * s), (cx + 0.5 * s, cy - 0.5 * s),
                (cx, cy - 0.85 * s), (cx - 0.5 * s, cy - 0.5 * s), (cx - 0.55 * s, cy + 0.1 * s)]
        ax.add_patch(Polygon(body, closed=True, facecolor=color, edgecolor="none", zorder=8))
        ip = [(cx, cy + 0.5 * s), (cx + 0.27 * s, cy - 0.05 * s),
              (cx, cy - 0.45 * s), (cx - 0.27 * s, cy - 0.05 * s)]
        ax.add_patch(Polygon(ip, closed=True, facecolor=inner, edgecolor="none", zorder=9))

    def grid_icon(cx, cy, s, color=FZ_EDGE):
        n, cell = 3, s / 3
        for i in range(n):
            for j in range(n):
                ax.add_patch(Rectangle((cx - s / 2 + i * cell, cy - s / 2 + j * cell),
                                       cell * 0.78, cell * 0.78, facecolor=color,
                                       edgecolor="white", lw=0.4, zorder=7))

    def dots_icon(cx, cy, s, color=TR_EDGE):
        for dx in (-0.6, 0.6):
            for dy in (-0.6, 0.6):
                ax.add_patch(Circle((cx + dx * s, cy + dy * s), s * 0.34,
                                    facecolor=color, edgecolor="none", zorder=7))

    def bubble(cx, cy, w, h, color="#7D3C98"):
        ax.add_patch(FancyBboxPatch((cx - w / 2, cy - h / 2), w, h,
                     boxstyle="round,pad=0,rounding_size=0.8", facecolor="white",
                     edgecolor=color, lw=1.2, zorder=7))
        ax.add_patch(Polygon([(cx - 0.18 * w, cy - h / 2), (cx + 0.02 * w, cy - h / 2),
                              (cx - 0.16 * w, cy - h)], closed=True, facecolor="white",
                             edgecolor=color, lw=1.0, zorder=7))
        for k in (-1, 0, 1):
            ax.add_patch(Circle((cx + k * 0.26 * w, cy), 0.07 * w, facecolor=color, zorder=8))

    def doc_icon(cx, cy, w, h, color=OUT_EDGE):
        ax.add_patch(Rectangle((cx - w / 2, cy - h / 2), w, h, facecolor="white",
                               edgecolor=color, lw=1.2, zorder=7))
        ax.add_patch(Polygon([(cx + w / 2 - 0.32 * w, cy + h / 2), (cx + w / 2, cy + h / 2),
                              (cx + w / 2, cy + h / 2 - 0.32 * h)], closed=True,
                             facecolor=color, edgecolor="none", zorder=8))
        for k in range(3):
            yy = cy + h * 0.18 - k * h * 0.22
            ax.plot([cx - w / 3, cx + w / 5], [yy, yy], color=color, lw=1.0, zorder=8)

    def chip(cx, cy, w, h, text, color):
        ax.add_patch(FancyBboxPatch((cx - w / 2, cy - h / 2), w, h,
                     boxstyle="round,pad=0,rounding_size=0.7", facecolor=color,
                     edgecolor="none", zorder=7))
        ax.text(cx, cy, text, ha="center", va="center", fontsize=6.3,
                color="white", zorder=8)

    def rbox(cx, cy, w, h, fill, edge):
        ax.add_patch(FancyBboxPatch((cx - w / 2, cy - h / 2), w, h,
                     boxstyle="round,pad=0,rounding_size=1.8", facecolor=fill,
                     edgecolor=edge, lw=1.8, zorder=5))

    def marker(cx, cy, w, h, kind):
        mx, my = cx - w / 2 + 2.8, cy + h / 2 - 2.8
        if kind == "frozen":
            snow(mx, my, 2.0)
        else:
            flame(mx, my, 2.3)

    yrow = 24

    # ---- input: anh X-quang that (hoac icon neu khong co) ----
    x_img_left, img_right = 3, 22
    used_image = False
    try:
        if CXR_IMAGE and Path(CXR_IMAGE).exists():
            im = plt.imread(CXR_IMAGE)
            ih, iw = im.shape[0], im.shape[1]
            th = 22
            tw = th * iw / ih
            ax.imshow(im, extent=[x_img_left, x_img_left + tw, yrow - th / 2, yrow + th / 2],
                      cmap="gray" if im.ndim == 2 else None, zorder=4)
            ax.add_patch(FancyBboxPatch((x_img_left, yrow - th / 2), tw, th,
                         boxstyle="round,pad=0,rounding_size=0.4", facecolor="none",
                         edgecolor="#0D2236", lw=2, zorder=5))
            img_right = x_img_left + tw
            used_image = True
    except Exception as e:  # noqa: BLE001
        print("  (khong doc duoc anh, dung icon):", e)
    if not used_image:
        cx = 12
        ax.add_patch(FancyBboxPatch((cx - 9, yrow - 9), 18, 18,
                     boxstyle="round,pad=0,rounding_size=1.4", facecolor="#16324F",
                     edgecolor="#0D2236", lw=1.8, zorder=4))
        for sgn in (-0.17, 0.17):
            ax.add_patch(Ellipse((cx + sgn * 18, yrow), 18 * 0.24, 18 * 0.5,
                                 facecolor="#CFE0EE", edgecolor="none", zorder=5))
        img_right = cx + 9
    ax.text((x_img_left + img_right) / 2, yrow - 13.5, "Chest X-ray\n(input image)",
            ha="center", va="top", fontsize=8)

    # ---- cac module chinh ----
    enc = (42, 22, 18)     # cx, w, h
    prj = (74, 22, 18)
    llm = (114, 22, 18)
    out = (140, 16, 20)

    rbox(enc[0], yrow, enc[1], enc[2], FZ_FILL, FZ_EDGE)
    grid_icon(enc[0], yrow + 4.2, 4.5)
    ax.text(enc[0], yrow - 3.5, "RAD-DINO\n(ViT-B/14)", ha="center", va="center", fontsize=8.2, zorder=10)
    marker(enc[0], yrow, enc[1], enc[2], "frozen")

    rbox(prj[0], yrow, prj[1], prj[2], TR_FILL, TR_EDGE)
    dots_icon(prj[0], yrow + 4.2, 2.3)
    ax.text(prj[0], yrow - 3.5, "MLP Projection\n(32 queries)", ha="center", va="center", fontsize=8.2, zorder=10)
    marker(prj[0], yrow, prj[1], prj[2], "train")

    rbox(llm[0], yrow, llm[1], llm[2], TR_FILL, TR_EDGE)
    bubble(llm[0], yrow + 4.4, 9, 4.6)
    ax.text(llm[0], yrow - 3.8, "Vicuna-7B\n+ LoRA", ha="center", va="center", fontsize=8.2, zorder=10)
    marker(llm[0], yrow, llm[1], llm[2], "train")

    rbox(out[0], yrow, out[1], out[2], OUT_FILL, OUT_EDGE)
    doc_icon(out[0], yrow + 5, 8, 7)
    ax.text(out[0], yrow - 4.5, "Findings /\nImpression /\nAnswer", ha="center", va="center", fontsize=7.4, zorder=10)

    # ---- mui ten + dai visual tokens ----
    arr(ax, img_right + 1, yrow, enc[0] - enc[1] / 2, yrow, color=ARROW, lw=2.2)
    arr(ax, enc[0] + enc[1] / 2, yrow, prj[0] - prj[1] / 2, yrow,
        "P  (1369×768)", color=ARROW, lw=2.2, fontsize=7.0)
    # visual tokens strip giua projection va LLM
    ts_x0, ts_x1 = prj[0] + prj[1] / 2 + 2, llm[0] - llm[1] / 2 - 2
    arr(ax, prj[0] + prj[1] / 2, yrow, llm[0] - llm[1] / 2, yrow, color=ARROW, lw=2.2)
    n = 8
    sq = (ts_x1 - ts_x0) / (n * 1.5)
    for i in range(n):
        ax.add_patch(Rectangle((ts_x0 + i * sq * 1.5, yrow - sq / 2), sq, sq,
                               facecolor="white", edgecolor=TR_EDGE, lw=0.9, zorder=6))
    ax.text((ts_x0 + ts_x1) / 2, yrow + 4.5, "32 visual tokens\nV (32×4096)",
            ha="center", va="bottom", fontsize=6.8, color=ARROW)
    arr(ax, llm[0] + llm[1] / 2, yrow, out[0] - out[1] / 2, yrow, color=ARROW, lw=2.2)

    # ---- nhanh CheXpert (PNU) ----
    cx_c, cy_c, w_c, h_c = 74, 52, 36, 15
    rbox(cx_c, cy_c, w_c, h_c, FZ_FILL, FZ_EDGE)
    marker(cx_c, cy_c, w_c, h_c, "frozen")
    ax.text(cx_c, cy_c + 4, "CheXpert Classifier  (14 pathologies)",
            ha="center", va="center", fontsize=7.6, zorder=10)
    chip(cx_c - 11, cy_c - 2.5, 11, 4.2, "Positive", POS)
    chip(cx_c, cy_c - 2.5, 11, 4.2, "Negative", NEG)
    chip(cx_c + 11, cy_c - 2.5, 11, 4.2, "Uncertain", UNC)
    arr(ax, enc[0], yrow + enc[2] / 2, cx_c - w_c / 2 + 4, cy_c - h_c / 2, color=ARROW, lw=1.8)
    arr(ax, cx_c + w_c / 2 - 4, cy_c - h_c / 2, llm[0], yrow + llm[2] / 2,
        "PNU string", dashed=True, color=ARROW, lw=1.8, fontsize=7.0)

    # ---- legend (icon) ----
    snow(98, 5, 2.2); ax.text(101, 5, "Frozen", va="center", fontsize=8)
    flame(122, 5, 2.6); ax.text(125, 5, "Trainable", va="center", fontsize=8)

    ax.set_xlim(0, W); ax.set_ylim(0, H); ax.set_aspect("equal")
    fig.savefig(OUT / "fig01_pipeline.png", dpi=180, bbox_inches="tight", facecolor="white")
    fig.savefig(OUT / "fig01_pipeline.svg", bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print("[+] fig01_pipeline.png + .svg", "(anh that)" if used_image else "(icon)")


# ---- cac hinh con lai (don gian, den trang) --------------------------------

def fig02_data_pipeline():
    fig, ax = canvas(165, 24)
    hrow(ax, 12, [
        (16, 22, "Data Sources", False),
        (49, 22, "[Local]\nSelection", False),
        (82, 26, "[Cloud VM]\nImage download", False),
        (117, 22, "Resize + shard", False),
        (151, 18, "Unified\nJSON", False),
    ], h=14)
    save(fig, "fig02_data_pipeline.png")


def fig03_selection_funnel():
    fig, ax = canvas(70, 86)
    steps = [
        (62, "~227k studies"),
        (54, "(a) Frontal only (PA/AP)"),
        (46, "(b) Findings + Impression"),
        (40, "(c) Length filter"),
        (34, "(d) Stratified sampling"),
        (44, "50k  (40k / 5k / 5k)"),
    ]
    ys = list(range(78, 6, -13))
    prev = None
    for (w, text), y in zip(steps, ys):
        box(ax, 35, y, w, 9, text)
        if prev is not None:
            arr(ax, 35, prev - 4.5, 35, y + 4.5)
        prev = y
    save(fig, "fig03_selection_funnel.png")


def fig04_sample_creation():
    fig, ax = canvas(96, 50)
    box(ax, 16, 25, 18, 10, "1 study")
    box(ax, 44, 25, 20, 10, "1 frontal\nimage")
    arr(ax, 25, 25, 34, 25)
    box(ax, 80, 38, 30, 12, "findings sample\nimage -> findings")
    box(ax, 80, 12, 34, 12, "impression sample\nimage + GT findings\n-> impression")
    arr(ax, 54, 27, 65, 38)
    arr(ax, 54, 23, 63, 14)
    arr(ax, 80, 32, 80, 18, "cascade")
    save(fig, "fig04_sample_creation.png")


def fig05_patchify():
    """RAD-DINO (ViT-B/14) ve lai theo kien truc goc — co patch grid, cot embedding
    neuron, va khoi Transformer Encoder mo rong. Tat ca FROZEN."""
    FZ = "#2E86C1"
    fig, ax = canvas(176, 70)

    # (a) anh dau vao
    ax.add_patch(FancyBboxPatch((4, 28), 16, 16, boxstyle="round,pad=0,rounding_size=1.2",
                 facecolor="#16324F", edgecolor="#0D2236", lw=1.6, zorder=4))
    for s in (-0.17, 0.17):
        ax.add_patch(Ellipse((12 + s * 16, 36), 4, 8, facecolor="#CFE0EE",
                             edgecolor="none", zorder=5))
    ax.text(12, 25, "Chest X-ray\n518×518", ha="center", va="top", fontsize=7.4)

    # (b) patch grid 6x6
    gx, gy, gs = 28, 28, 16
    cell = gs / 6
    for i in range(6):
        for j in range(6):
            ax.add_patch(Rectangle((gx + i * cell, gy + j * cell), cell * 0.86,
                         cell * 0.86, facecolor="#D6EAF8", edgecolor=FZ, lw=0.5, zorder=5))
    ax.text(gx + gs / 2, 25, "1369 patches\n(14×14 px)", ha="center", va="top", fontsize=7.4)
    arr(ax, 20, 36, 27.5, 36, color="#5D6D7E", lw=1.6)

    # (c) linear projection -> patch+pos embedding (cot neuron)
    box(ax, 56, 36, 14, 9, "Linear\nProjection")
    arr(ax, gx + gs, 36, 49, 36, color="#5D6D7E", lw=1.6)
    emb = ncol(ax, 76, 36, 7, edge=FZ)
    arr(ax, 63, 36, 73.5, 36, color="#5D6D7E", lw=1.6)
    ax.text(76, 20, "patch + position\nembedding (768-d)", ha="center", va="top", fontsize=7.2)

    # (d) Transformer Encoder block (mo rong) x12
    x0, x1 = 92, 150
    rrect(ax, x0, 12, x1, 58, fill="#F4F9FD", edge=FZ, lw=1.6)
    ax.text((x0 + x1) / 2, 61, "Transformer Encoder  × 12   (frozen)",
            ha="center", fontsize=7.8, weight="bold")
    rows = [(52, "LayerNorm"), (45, "Multi-Head Self-Attention"),
            (38, "⊕  residual"), (31, "LayerNorm"),
            (24, "MLP  (GELU, 4×)"), (17, "⊕  residual")]
    for cy, label in rows:
        box(ax, (x0 + x1) / 2, cy, 50, 5.4, label)
    arr(ax, 79, 36, x0, 36, color="#5D6D7E", lw=1.6)

    # (e) output tokens (cot neuron, [CLS] highlight)
    ys = ncol(ax, 166, 36, 7, edge=FZ)
    ax.add_patch(Circle((166, ys[-1]), 1.0, facecolor="#FAD7A0", edgecolor="#E67E22",
                 lw=1.0, zorder=7))   # [CLS]
    arr(ax, x1, 36, 163.5, 36, color="#5D6D7E", lw=1.6)
    ax.text(166, 18, "[CLS] + 1369\npatch tokens", ha="center", va="top", fontsize=7.2)

    save(fig, "fig05_patchify.png")


def fig06_projection():
    """MLP Projection — phan MLP ve bang cot neuron (768 -> 1024 -> 4096)."""
    TR = "#E67E22"
    fig, ax = canvas(150, 50)

    box(ax, 15, 33, 22, 10, "Patch features\n1369 × 768", frozen=True)
    box(ax, 15, 14, 20, 8, "32 query\ntokens")
    box(ax, 45, 24, 22, 12, "Cross-Attention\n(8 heads)")
    arr(ax, 26, 33, 33.5, 27, color="#5D6D7E", lw=1.5)
    arr(ax, 25, 14, 33.5, 21, color="#5D6D7E", lw=1.5)
    ax.text(45, 16, "→ 32 × 768", ha="center", fontsize=6.8)

    # MLP per-token: 768 -> 1024 -> 4096 (so vong tron rut gon)
    xa, xb, xc = 76, 100, 124
    ax.text(100, 47, "per-token MLP  (applied to each of the 32 tokens)",
            ha="center", fontsize=7.0, style="italic")
    ya = ncol(ax, xa, 24, 5, edge=TR)
    yb = ncol(ax, xb, 24, 7, edge=TR)
    yc = ncol(ax, xc, 24, 6, edge=TR)
    nconnect(ax, xa, ya, xb, yb)
    nconnect(ax, xb, yb, xc, yc)
    # ve lai vong tron de noi nam tren day
    ncol(ax, xa, 24, 5, edge=TR); ncol(ax, xb, 24, 7, edge=TR); ncol(ax, xc, 24, 6, edge=TR)
    arr(ax, 56, 24, xa - 4, 24, color="#5D6D7E", lw=1.5)
    ax.text(xa, 11, "768", ha="center", fontsize=7)
    ax.text(xb, 11, "1024\n(tap)", ha="center", fontsize=7)
    ax.text(xc, 11, "4096", ha="center", fontsize=7)
    ax.text(xa, 33, "GELU+Dropout", ha="center", fontsize=6.3, color="#888")

    # nhanh ITC head tu lop 1024
    box(ax, xb, 42, 20, 7, "ITC head (Stage 1)", dashed=True)
    arr(ax, xb, 24 + 7 * 2.9 / 2 + 1, xb, 38.5, color="#5D6D7E", lw=1.3, dashed=True)
    # output
    box(ax, 142, 24, 14, 10, "V\n32×4096")
    arr(ax, xc + 4, 24, 135, 24, color="#5D6D7E", lw=1.5)
    save(fig, "fig06_projection.png")


def fig07_chexpert_pnu():
    fig, ax = canvas(130, 22)
    hrow(ax, 11, [
        (15, 20, "RAD-DINO\n[CLS]", True),
        (43, 16, "MLP head", False),
        (68, 16, "14 x 3\nlogits", False),
        (93, 18, "format_pnu()", False),
        (118, 18, "PNU string", False),
    ], h=12)
    save(fig, "fig07_chexpert_pnu.png")


def fig08_lora():
    """Vicuna-7B decoder (frozen, 4-bit) + LoRA (trainable) ve lai theo kien truc goc.
    FFN ve bang cot neuron. Xanh = frozen, cam = trainable."""
    FZ, TR = "#2E86C1", "#E67E22"
    fig, ax = canvas(178, 74)

    # (a) input tokens (cot neuron)
    yi = ncol(ax, 12, 37, 8, edge="#7D7D7D")
    ax.text(12, 18, "input tokens\n(text + 32 visual)", ha="center", va="top", fontsize=7)
    box(ax, 32, 37, 16, 10, "Token + Pos\nEmbedding", frozen=True)
    arr(ax, 15, 37, 24, 37, color="#5D6D7E", lw=1.5)

    # (b) Decoder block x32
    x0, x1 = 46, 138
    rrect(ax, x0, 10, x1, 64, fill="#F7FAFC", edge="#566573", lw=1.6)
    ax.text((x0 + x1) / 2, 67, "Decoder block  × 32", ha="center", fontsize=7.8, weight="bold")
    arr(ax, 40, 37, x0, 37, color="#5D6D7E", lw=1.5)

    # masked self-attention (frozen) + LoRA (trainable)
    box(ax, 70, 54, 38, 9, "Masked Multi-Head Self-Attention\nq, k, v, o  (4-bit, frozen)", frozen=True)
    # LoRA motif: 2 cot neuron nho (A -> B) mau cam
    la = ncol(ax, 104, 54, 3, r=0.8, vgap=2.2, edge=TR)
    lb = ncol(ax, 110, 54, 2, r=0.8, vgap=2.2, edge=TR)
    nconnect(ax, 104, la, 110, lb, color="#F0C9A0", lw=0.5)
    ncol(ax, 104, 54, 3, r=0.8, vgap=2.2, edge=TR); ncol(ax, 110, 54, 2, r=0.8, vgap=2.2, edge=TR)
    ax.text(117, 54, "LoRA A→B\n(r=16, trainable)", ha="left", va="center", fontsize=6.4, color=TR)
    ax.text(96, 49, "⊕", ha="center", fontsize=9)

    box(ax, 70, 43, 38, 5, "⊕  +  LayerNorm")

    # feed-forward (frozen) ve bang cot neuron
    ax.text(70, 37, "Feed-Forward MLP  (frozen)", ha="center", fontsize=6.8)
    fa = ncol(ax, 58, 28, 4, r=0.9, edge=FZ)
    fb = ncol(ax, 70, 28, 6, r=0.9, edge=FZ)
    fc = ncol(ax, 82, 28, 4, r=0.9, edge=FZ)
    nconnect(ax, 58, fa, 70, fb); nconnect(ax, 70, fb, 82, fc)
    ncol(ax, 58, 28, 4, r=0.9, edge=FZ); ncol(ax, 70, 28, 6, r=0.9, edge=FZ); ncol(ax, 82, 28, 4, r=0.9, edge=FZ)
    box(ax, 70, 15, 38, 5, "⊕  +  LayerNorm")

    # (c) LM head -> output
    box(ax, 150, 37, 16, 10, "LM Head", frozen=True)
    arr(ax, x1, 37, 142, 37, color="#5D6D7E", lw=1.5)
    box(ax, 170, 37, 14, 12, "next\ntoken")
    arr(ax, 158, 37, 163, 37, color="#5D6D7E", lw=1.5)

    # legend
    ax.add_patch(FancyBboxPatch((46, 2), 4, 3, boxstyle="round,pad=0,rounding_size=0.6",
                 facecolor="#D6EAF8", edgecolor=FZ, lw=1.2))
    ax.text(51, 3.5, "Frozen (4-bit base)", va="center", fontsize=7)
    ax.add_patch(FancyBboxPatch((96, 2), 4, 3, boxstyle="round,pad=0,rounding_size=0.6",
                 facecolor="#FDEBD0", edgecolor=TR, lw=1.2))
    ax.text(101, 3.5, "Trainable (LoRA)", va="center", fontsize=7)
    save(fig, "fig08_lora.png")


def fig09_prompt():
    fig, ax = canvas(124, 40)
    segs = [("[SYSTEM]", 18), ("<image>", 14), ("PNU / Findings", 24),
            ("instruction", 22), ("ASSISTANT: target", 28)]
    x = 4
    for text, w in segs:
        box(ax, x + w / 2, 30, w, 9, text)
        x += w + 1
    ax.text(4, 37, "Prompt anatomy", fontsize=8, weight="bold")
    box(ax, 22, 10, 28, 9, "<image>  (1 token)")
    box(ax, 78, 10, 30, 9, "32 visual tokens")
    arr(ax, 36, 10, 63, 10, "expand")
    ax.text(104, 10, "mask/labels +31\nnon-target = -100",
            ha="left", va="center", fontsize=7.3)
    ax.text(4, 17, "Image-token expansion", fontsize=8, weight="bold")
    save(fig, "fig09_prompt.png")


def fig10_curriculum():
    fig, ax = canvas(120, 28)
    box(ax, 20, 14, 30, 14, "Stage 0\nCheXpert head")
    box(ax, 60, 14, 32, 14, "Stage 1\nProjection + ITC")
    box(ax, 102, 14, 30, 14, "Stage 2\nProjection + LoRA")
    arr(ax, 35, 14, 44, 14)
    arr(ax, 76, 14, 87, 14, "projection\nweights")
    save(fig, "fig10_curriculum.png")


def fig11_contrastive():
    fig, ax = canvas(120, 36)
    box(ax, 18, 27, 30, 10, "image -> RAD-DINO\n-> proj -> ITC")
    box(ax, 60, 27, 22, 10, "128-d (image)")
    arr(ax, 33, 27, 49, 27)
    box(ax, 18, 8, 32, 10, "findings -> CXR-BERT\n(offline cache)")
    box(ax, 60, 8, 22, 10, "128-d (text)")
    arr(ax, 34, 8, 49, 8)
    box(ax, 100, 17, 30, 14, "similarity B x B\nInfoNCE")
    arr(ax, 71, 27, 86, 20)
    arr(ax, 71, 8, 86, 14)
    save(fig, "fig11_contrastive.png")


def fig12_loss_mask():
    fig, ax = canvas(124, 20)
    segs = [("[SYSTEM]", 18, True), ("32 visual", 18, True), ("PNU", 12, True),
            ("instruction", 24, True), ("pad", 12, True), ("response", 28, False)]
    x = 4
    for text, w, ignored in segs:
        box(ax, x + w / 2, 12, w, 9, text, frozen=ignored)
        x += w
    ax.text(46, 4, "-100  (ignored by loss)", ha="center", fontsize=7.5)
    ax.text(110, 4, "loss computed", ha="center", fontsize=7.5)
    save(fig, "fig12_loss_mask.png")


def fig13_eval_flow():
    fig, ax = canvas(120, 46)
    box(ax, 14, 30, 18, 10, "test image")
    box(ax, 44, 30, 26, 12, "CheXpert PNU\n(+GT findings\nif impression)")
    box(ax, 78, 30, 20, 10, "greedy\ndecode")
    box(ax, 106, 30, 18, 10, "hypothesis")
    arr(ax, 23, 30, 31, 30)
    arr(ax, 57, 30, 68, 30)
    arr(ax, 88, 30, 97, 30)
    box(ax, 40, 8, 24, 9, "NLG metrics")
    box(ax, 72, 8, 24, 9, "CheXbert F1")
    box(ax, 104, 8, 28, 9, "VQA: EM / F1 / judge")
    arr(ax, 106, 25, 52, 13)
    arr(ax, 106, 25, 78, 13)
    arr(ax, 106, 25, 104, 13)
    save(fig, "fig13_eval_flow.png")


def fig14_chexbert_f1():
    fig, ax = canvas(116, 36)
    box(ax, 16, 27, 26, 10, "generated report")
    box(ax, 16, 9, 26, 10, "reference report")
    box(ax, 56, 18, 24, 11, "CheXbert\nlabeler")
    box(ax, 90, 18, 24, 11, "14-label\nvectors")
    arr(ax, 29, 27, 44, 21)
    arr(ax, 29, 9, 44, 15)
    arr(ax, 68, 18, 78, 18)
    ax.text(90, 8, "macro-F1 / P / R", ha="center", fontsize=7.8)
    save(fig, "fig14_chexbert_f1.png")


def fig15_workflow():
    """So do workflow tong the cua du an (kieu flowchart pastel, doc tu tren xuong).
    Chinh mau o dict C, chinh chu/vi tri o cac lenh fb(...) ben duoi."""
    C = {"blue": ("#D6EAF8", "#2E86C1"), "purple": ("#E8DAEF", "#8E44AD"),
         "orange": ("#FDEBD0", "#E67E22"), "red": ("#FADBD8", "#C0392B"),
         "green": ("#D5F5E3", "#27AE60"), "teal": ("#D1F2EB", "#16A085"),
         "yellow": ("#FCF3CF", "#D4AC0D"), "gray": ("#D5D8DC", "#566573")}
    EDGE = "#566573"
    W, H = 136, 176
    fig, ax = canvas(W, H)

    def fb(cx, cy, w, h, text, color, fs=7.2):
        fill, edge = C[color]
        ax.add_patch(FancyBboxPatch((cx - w / 2, cy - h / 2), w, h,
                     boxstyle="round,pad=0,rounding_size=1.3", facecolor=fill,
                     edgecolor=edge, lw=1.4, zorder=5))
        ax.text(cx, cy, text, ha="center", va="center", fontsize=fs, zorder=6)
        return (cx, cy, w, h)

    def group(x0, y0, x1, y1, title):
        ax.add_patch(FancyBboxPatch((x0, y0), x1 - x0, y1 - y0,
                     boxstyle="round,pad=0,rounding_size=2", facecolor="#FBFCFC",
                     edgecolor=EDGE, lw=1.3, zorder=1))
        ax.text(x0 + 2.5, y1 - 3, title, ha="left", va="center", fontsize=7.6,
                weight="bold", zorder=2)

    def dn(a, b, text=None, dashed=False):
        arr(ax, a[0], a[1] - a[3] / 2, b[0], b[1] + b[3] / 2, text=text,
            dashed=dashed, color=EDGE, lw=1.3)

    # ---- top: data + split + branches ----
    data = fb(68, 169, 56, 9, "MIMIC-CXR  (50,000 frontal images)", "blue", 8)
    split = fb(68, 157, 50, 8, "Stratified patient-disjoint split", "purple", 7.4)
    dn(data, split)
    tr = fb(34, 145, 26, 8, "Train (~80%)", "blue", 7.2)
    va = fb(68, 145, 22, 8, "Val (~10%)", "blue", 7.2)
    te = fb(102, 145, 24, 8, "Test (~10%)", "blue", 7.2)
    for b in (tr, va, te):
        arr(ax, split[0], split[1] - 4, b[0], b[1] + 4, color=EDGE, lw=1.3)

    # ---- groups ----
    group(8, 48, 64, 139, "Training Pipeline (Stage 0 / 1 / 2)")
    group(72, 67, 128, 123, "Inference Pipeline")

    # ---- training stack ----
    t1 = fb(36, 130, 51, 8, "Image preprocessing (resize 518, center-crop)", "orange")
    t2 = fb(36, 118, 51, 10, "RAD-DINO backbone — frozen\n(ViT-B/14, DINOv2 on CXR)", "red")
    t3 = fb(36, 106, 51, 8, "MLP Projection head (32 visual tokens)", "blue")
    t4 = fb(36, 93, 51, 9, "Stage 0 — CheXpert classifier head\n(BCE, U-MultiClass)", "green")
    t5 = fb(36, 80, 51, 9, "Stage 1 — Contrastive alignment\n(ITC / InfoNCE · LLM not loaded)", "green")
    t6 = fb(36, 68, 51, 9, "Stage 2 — QLoRA instruction tuning\n(causal LM loss · q,k,v,o · 4-bit)", "orange")
    t7 = fb(36, 57, 44, 8, "Best Checkpoint (lowest val loss)", "blue")
    for a, b in [(t1, t2), (t2, t3), (t3, t4), (t4, t5), (t5, t6), (t6, t7)]:
        dn(a, b)
    arr(ax, tr[0], tr[1] - 4, t1[0], t1[1] + 4, color=EDGE, lw=1.3)

    # ---- inference stack ----
    i1 = fb(100, 114, 50, 10, "CheXpert head → PNU\n(Positive / Negative / Uncertain)", "yellow")
    i2 = fb(100, 101, 50, 10, "Prompt assembly\n(image + PNU + instruction)", "purple")
    i3 = fb(100, 89, 50, 8, "Greedy decoding", "yellow")
    i4 = fb(100, 77, 50, 10, "Output:\nFindings / Impression / VQA", "red")
    for a, b in [(i1, i2), (i2, i3), (i3, i4)]:
        dn(a, b)
    arr(ax, te[0], te[1] - 4, i1[0], i1[1] + 5, color=EDGE, lw=1.3)
    arr(ax, t7[0] + t7[2] / 2, t7[1], i4[0] - i4[2] / 2, i4[1], "trained model",
        color=EDGE, lw=1.3, fontsize=6.6)

    # ---- val feedback (early stopping): di doc khe giua 2 nhom ----
    ax.plot([68, 68], [va[1] - 4, 59], ls="--", color=EDGE, lw=1.2, zorder=2)
    arr(ax, 68, 59, t7[0] + t7[2] / 2, t7[1] + 1, dashed=True, color=EDGE, lw=1.2)
    ax.text(69.5, 100, "validation / early stopping", rotation=90, va="center",
            ha="left", fontsize=6.4, color=EDGE)

    # ---- bottom row ----
    ev = fb(100, 35, 54, 14, "Evaluation Metrics\nBLEU · ROUGE · METEOR · BERTScore\nCheXbert F1 · VQA EM / token-F1", "gray", 6.4)
    dn(i4, ev)

    fig.savefig(OUT / "fig15_workflow.png", dpi=180, bbox_inches="tight", facecolor="white")
    fig.savefig(OUT / "fig15_workflow.svg", bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print("[+] fig15_workflow.png + .svg")


# ---- EDA charts redrawn in English (overwrite the notebook-extracted ones) ----
import numpy as np


def eda_chexpert_labels():
    labels = ["No Finding", "Support Devices", "Pleural Effusion", "Lung Opacity",
              "Atelectasis", "Cardiomegaly", "Edema", "Pneumonia", "Consolidation",
              "Pneumothorax", "Enlarged Cardiomediastinum", "Lung Lesion", "Fracture",
              "Pleural Other"]
    pos = [75455, 66558, 54300, 51525, 45808, 44845, 27018, 16556, 10778, 10358, 7179, 6284, 4390, 2011]
    unc = [0, 237, 5814, 3831, 10327, 6043, 13174, 18291, 4331, 1134, 9375, 1141, 555, 765]
    neg = [0, 3486, 27158, 3069, 1531, 15911, 25641, 24338, 7967, 42356, 5283, 862, 886, 126]
    x = np.arange(len(labels)); w = 0.27
    fig, ax = plt.subplots(figsize=(13, 5))
    ax.bar(x - w, pos, w, label="Positive", color="#E74C3C")
    ax.bar(x, unc, w, label="Uncertain", color="#F39C12")
    ax.bar(x + w, neg, w, label="Negative", color="#2ECC71")
    ax.set_xticks(x); ax.set_xticklabels(labels, rotation=40, ha="right", fontsize=9)
    ax.set_ylabel("Number of studies")
    ax.set_title("CheXpert labels — Positive / Uncertain / Negative (full dataset)")
    ax.legend(); fig.tight_layout()
    fig.savefig(OUT / "eda_chexpert_labels.png", dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig); print("[+] eda_chexpert_labels.png (EN)")


def eda_views():
    views = ["AP", "PA", "LATERAL", "LL", "Unknown"]
    cnt = [147173, 96161, 82853, 35133, 15769]
    colors = list(plt.cm.Set2.colors[:len(views)])
    fig, (a1, a2) = plt.subplots(1, 2, figsize=(13, 4))
    bars = a1.bar(views, cnt, color=colors)
    a1.bar_label(bars, fmt="%d", fontsize=8)
    a1.set_ylabel("Number of images"); a1.set_title("Images by view position")
    a2.pie(cnt, labels=views, autopct="%1.1f%%", colors=colors)
    a2.set_title("View-position share")
    fig.suptitle("View position distribution — full dataset")
    fig.tight_layout()
    fig.savefig(OUT / "eda_views.png", dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig); print("[+] eda_views.png (EN)")


def eda_imgs_per_study():
    d = {1: 102675, 2: 103481, 3: 19442, 4: 2097, 5: 99, 6: 33, 7: 2, 8: 4, 9: 1, 11: 1}
    ks = list(d); vs = [d[k] for k in ks]
    fig, ax = plt.subplots(figsize=(9, 4))
    ax.bar([str(k) for k in ks], vs, color="#4C72B0")
    for k, v in zip(ks, vs):
        ax.text(str(k), v + 1500, f"{v:,}", ha="center", va="bottom", fontsize=8)
    ax.set_xlabel("Images per study"); ax.set_ylabel("Number of studies")
    ax.set_title("Images per study — full dataset (mean 1.66, max 11)")
    fig.tight_layout()
    fig.savefig(OUT / "eda_imgs_per_study.png", dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig); print("[+] eda_imgs_per_study.png (EN)")


def eda_report_length():
    secs = ["Findings", "Impression"]; med = [45, 16]; mean = [48.7, 24.6]
    x = np.arange(2); w = 0.35
    fig, ax = plt.subplots(figsize=(7, 4))
    ax.bar(x - w / 2, med, w, label="Median", color="#4C72B0")
    ax.bar(x + w / 2, mean, w, label="Mean", color="#DD8452")
    for i, (m, a) in enumerate(zip(med, mean)):
        ax.text(i - w / 2, m + 0.6, str(m), ha="center", fontsize=8)
        ax.text(i + w / 2, a + 0.6, str(a), ha="center", fontsize=8)
    ax.set_xticks(x); ax.set_xticklabels(secs); ax.set_ylabel("Word count")
    ax.set_title("Report section length (words) — full dataset")
    ax.legend(); fig.tight_layout()
    fig.savefig(OUT / "eda_report_length.png", dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig); print("[+] eda_report_length.png (EN)")


def main():
    for f in [fig01_pipeline, fig02_data_pipeline, fig03_selection_funnel,
              fig04_sample_creation, fig05_patchify, fig06_projection,
              fig07_chexpert_pnu, fig08_lora, fig09_prompt, fig10_curriculum,
              fig11_contrastive, fig12_loss_mask, fig13_eval_flow,
              fig14_chexbert_f1, fig15_workflow,
              eda_chexpert_labels, eda_views, eda_imgs_per_study, eda_report_length]:
        f()
    print("\nXong. Anh o:", OUT)


if __name__ == "__main__":
    main()