"""
Trích các biểu đồ EDA đã render sẵn trong output cell của notebook -> file PNG.
Chay: python export_eda_figs.py   (anh ra docs/figures/)

Neu ban chay lai notebook va cell index doi, sua dict PICKS ben duoi.
"""
import base64, json
from pathlib import Path

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

# (notebook, cell_index) -> ten file png
PICKS = {
    ("eda_full.ipynb", 25): "eda_chexpert_labels.png",   # P/U/N 14 nhan
    ("eda_full.ipynb", 17): "eda_views.png",             # view position bar+pie
    ("eda_full.ipynb", 14): "eda_imgs_per_study.png",    # so anh / study
    ("eda_full.ipynb", 31): "eda_report_length.png",     # do dai findings/impression
    ("eda_full.ipynb", 40): "eda_vqa_types.png",         # semantic + content type
    ("build_subset_local.ipynb", 16): "eda_prevalence_compare.png",  # full vs elig vs subset
    ("build_subset_local.ipynb", 17): "eda_vqa_compare.png",         # VQA full vs subset
}


def first_png(cell):
    for o in cell.get("outputs", []):
        img = (o.get("data", {}) or {}).get("image/png")
        if img:
            return img if isinstance(img, str) else "".join(img)
    return None


def main():
    for (nb_name, idx), out_name in PICKS.items():
        nb = json.load(open(NB_DIR / nb_name, encoding="utf-8"))
        if idx >= len(nb["cells"]):
            print(f"[x] {nb_name}: khong co cell {idx}")
            continue
        png_b64 = first_png(nb["cells"][idx])
        if not png_b64:
            print(f"[x] {nb_name} cell {idx}: khong co anh output")
            continue
        (OUT / out_name).write_bytes(base64.b64decode(png_b64))
        src = "".join(nb["cells"][idx].get("source", []))
        print(f"[+] {out_name:<28s} <- {nb_name} cell {idx}  ({src[:55].strip()}...)")


if __name__ == "__main__":
    main()