"""
Build scripts/cxrvlm_colab_inference.ipynb from cell sources defined here.

Run:
    python scripts/_build_inference_notebook.py
"""

import json
from pathlib import Path

NB_PATH = Path(__file__).parent / "cxrvlm_colab_inference.ipynb"


def md(cell_id, src):
    return {
        "cell_type": "markdown",
        "id": cell_id,
        "metadata": {},
        "source": src.splitlines(keepends=True),
    }


def code(cell_id, src):
    return {
        "cell_type": "code",
        "id": cell_id,
        "metadata": {},
        "execution_count": None,
        "outputs": [],
        "source": src.splitlines(keepends=True),
    }


CELLS = []

# ─────────────────────────────────────────────────────────────────────

CELLS.append(md("inf-0", """\
# CXR-VLM — Inference Notebook (Colab T4 / GCP)

Free-form inference on individual chest X-ray images. Use this when you want to:
- Run the trained model on a **new image** (not in the test split).
- Ask **custom VQA questions**.
- See the image + generated text side-by-side for qualitative checks.
- Iterate on prompts without re-running the heavy `evaluate.py` pipeline.

For aggregate test-set metrics use the companion notebook **`cxrvlm_colab_eval.ipynb`**.

### Sections
1. Selectors (which trained model to load + compute device)
2. Env + pip
3. Pull code from HF
4. Pull trained checkpoint + config snapshot from HF
5. Auto-detect GPU profile (bf16+FA2 on Ampere+/Ada, fp16+SDPA on Turing)
6. Build configs (4-bit Vicuna with the profile)
7. **Load model in-kernel** (one-time, ~5–8 min on T4 cold cache, ~3 min on L4)
8. **Inference helpers** (single image → findings / impression / report / VQA / full cascade)
9. **Examples** — single image, VQA, batch folder
"""))


CELLS.append(md("inf-select-md", """\
## 0. Select trained run + device
"""))


CELLS.append(code("inf-select", """\
# ── Platform ─────────────────────────────────────────────────────
PLATFORM = 'colab'          # 'kaggle' | 'colab' | 'lightning' | 'gcp' | 'local'

# ── Source repos on HuggingFace ──────────────────────────────────
HF_USER       = 'hieu3636'
HF_CODE_REPO  = f'{HF_USER}/cxr-vlm-code'
HF_RUNS_REPO  = f'{HF_USER}/cxr-vlm-runs'

# ── Which trained run to use ─────────────────────────────────────
RUN_ID        = 'MIMIC-CXR_resized_run_1'
CKPT_PICK     = 'best'                 # 'best' | 'last'

# ── Generation defaults (override per-call later if desired) ─────
MAX_NEW_TOKENS  = 300
TEMPERATURE     = 0.1
DO_SAMPLE       = False                # greedy by default — deterministic
NUM_BEAMS       = 1

assert PLATFORM  in ('kaggle', 'colab', 'lightning', 'gcp', 'local')
assert CKPT_PICK in ('best', 'last')

print(f'PLATFORM      = {PLATFORM}')
print(f'RUN_ID        = {RUN_ID}     (ckpt: stage2/{CKPT_PICK})')
"""))


CELLS.append(md("inf-env-md", "## 1. Environment + pip install"))


CELLS.append(code("inf-env", """\
import os
os.environ['CUDA_VISIBLE_DEVICES']            = '0'
os.environ['TOKENIZERS_PARALLELISM']          = 'false'
os.environ['BITSANDBYTES_NOWELCOME']          = '1'
os.environ['TRANSFORMERS_VERBOSITY']          = 'warning'
os.environ['PYTHONUNBUFFERED']                = '1'

import sys, shutil, subprocess
from pathlib import Path
"""))


CELLS.append(code("inf-pip", """\
# Same dependency window as the training/eval notebooks.
!pip uninstall -y -q torchao transformers bitsandbytes peft accelerate
!pip install -q -U bitsandbytes
!pip install -q \\
    'transformers>=4.46,<4.50' \\
    'peft>=0.13,<0.15' \\
    'accelerate>=1.0' \\
    'huggingface_hub>=0.27,<1.0' \\
    omegaconf sentencepiece 'protobuf>=3.20' \\
    pillow matplotlib

# Best-effort flash-attn install for Ampere+/Ada GPUs (L4, 3090, A10, A100, H100).
# Silent fail is OK — cxr_vlm.py auto-falls-back to SDPA if FA2 isn't importable.
import torch as _t
if _t.cuda.is_available() and _t.cuda.get_device_capability(0) >= (8, 0):
    print('[pip] Ampere+/Ada detected -> trying flash-attn (5-10 min if building from source)')
    !pip install -q flash-attn --no-build-isolation 2>&1 | tail -5
else:
    print('[pip] Pre-Ampere GPU (or no CUDA) -> skipping flash-attn install')
"""))


CELLS.append(code("inf-versions", """\
import torch, transformers, peft, huggingface_hub, httpx

print('torch          :', torch.__version__, '| cuda:', torch.cuda.is_available())
print('transformers   :', transformers.__version__)
print('peft           :', peft.__version__)
print('huggingface_hub:', huggingface_hub.__version__)
print('httpx          :', httpx.__version__)

# httpx 0.28+ shim — same as the other notebooks.
def _patch_httpx():
    if tuple(int(x) for x in httpx.__version__.split('.')[:2]) < (0, 28):
        return
    if getattr(httpx.Client, '_cxr_vlm_compat_patched', False):
        return
    def _make(orig):
        def patched(self, *args, **kwargs):
            if 'allow_redirects' in kwargs:
                kwargs['follow_redirects'] = kwargs.pop('allow_redirects')
            kwargs.pop('proxies', None)
            return orig(self, *args, **kwargs)
        return patched
    for cls in (httpx.Client, httpx.AsyncClient):
        for m in ('request', 'get', 'head', 'post', 'put',
                  'patch', 'delete', 'options'):
            if hasattr(cls, m):
                setattr(cls, m, _make(getattr(cls, m)))
    httpx.Client._cxr_vlm_compat_patched = True

_patch_httpx()

assert torch.cuda.is_available(), 'CUDA not available.'
_p = torch.cuda.get_device_properties(0)
print(f'GPU         : {_p.name}  ({_p.total_memory/1e9:.1f} GB)')
"""))


CELLS.append(md("inf-paths-md", "## 2. Paths + pull code"))


CELLS.append(code("inf-paths", """\
# WORK dir + HF_TOKEN bootstrap (platform-specific).
if PLATFORM == 'kaggle':
    from kaggle_secrets import UserSecretsClient
    os.environ['HF_TOKEN'] = UserSecretsClient().get_secret('HF_TOKEN')
    WORK = Path('/kaggle/working')
elif PLATFORM == 'colab':
    from google.colab import userdata
    os.environ['HF_TOKEN'] = userdata.get('HF_TOKEN')
    WORK = Path('/content')
elif PLATFORM == 'lightning':
    WORK = Path('/teamspace/studios/this_studio')
elif PLATFORM == 'gcp':
    # Vertex AI Workbench → /home/jupyter. Generic GCE VM → /workspace.
    # Pick whichever exists; create the latter if neither is writable.
    for _cand in (Path('/home/jupyter'), Path('/workspace')):
        if _cand.exists() or os.access(_cand.parent, os.W_OK):
            WORK = _cand
            break
    else:
        WORK = Path.home() / 'cxr-vlm-work'
else:  # 'local'
    WORK = Path.home() / 'cxr-vlm-work'
WORK.mkdir(parents=True, exist_ok=True)

assert os.environ.get('HF_TOKEN'), \\
    'HF_TOKEN missing — set it via the platform secrets UI before re-running.'

from huggingface_hub import snapshot_download

print(f'Pulling code from {HF_CODE_REPO} …')
CODE_SRC = Path(snapshot_download(
    repo_id   = HF_CODE_REPO,
    repo_type = 'model',
    token     = os.environ['HF_TOKEN'],
    local_dir = str(WORK / 'cxr-vlm-code'),
))

PROJECT = WORK / 'cxr_vlm'
if CODE_SRC.resolve() != PROJECT.resolve() and not PROJECT.exists():
    shutil.copytree(CODE_SRC, PROJECT)

os.chdir(PROJECT)
sys.path.insert(0, str(PROJECT))
print('PROJECT :', PROJECT)
print('WORK    :', WORK)
"""))


CELLS.append(md("inf-pull-run-md", """\
## 3. Pull trained checkpoint + config snapshot

Same as the eval notebook: pulls `{RUN_ID}/configs/` + `{RUN_ID}/stage2/{best|last}/` from `HF_RUNS_REPO`. No dataset payload needed for free-form inference.
"""))


CELLS.append(code("inf-pull-run", """\
from huggingface_hub import snapshot_download

RUN_PULL_ROOT = WORK / 'run_pull'
RUN_PULL_ROOT.mkdir(parents=True, exist_ok=True)

print(f'Pulling {RUN_ID}/{{configs,stage2/{CKPT_PICK}}} from {HF_RUNS_REPO} …')
snapshot_download(
    repo_id   = HF_RUNS_REPO,
    repo_type = 'model',
    token     = os.environ['HF_TOKEN'],
    allow_patterns = [
        f'{RUN_ID}/configs/**',
        f'{RUN_ID}/run_meta.json',
        f'{RUN_ID}/stage2/{CKPT_PICK}/**',
    ],
    local_dir = str(RUN_PULL_ROOT),
)

RUN_DIR_PULLED = RUN_PULL_ROOT / RUN_ID
CKPT_DIR_PULLED = RUN_DIR_PULLED / 'stage2' / CKPT_PICK
PROJ_PT         = CKPT_DIR_PULLED / 'checkpoint_projection.pt'
LORA_DIR        = CKPT_DIR_PULLED / 'checkpoint_lora'
CHEXPERT_PT     = CKPT_DIR_PULLED / 'checkpoint_chexpert_classifier.pt'

assert PROJ_PT.is_file(), f'Projection weights not found at {PROJ_PT}.'
assert (LORA_DIR / 'adapter_config.json').is_file(), \\
    f'LoRA adapter_config.json not found in {LORA_DIR}.'

print(f'  projection : {PROJ_PT}      ({PROJ_PT.stat().st_size/1e6:.1f} MB)')
print(f'  lora       : {LORA_DIR}/    ({sum(p.stat().st_size for p in LORA_DIR.rglob("*") if p.is_file())/1e6:.1f} MB)')
print(f'  chexpert   : {CHEXPERT_PT}  (exists: {CHEXPERT_PT.is_file()})')

SAVED_CFG_DIR   = RUN_DIR_PULLED / 'configs'
SAVED_TRAIN_CFG = SAVED_CFG_DIR / 'train_config.yaml'
SAVED_MODEL_CFG = SAVED_CFG_DIR / 'model_config.yaml'
"""))


CELLS.append(md("inf-gpu-md", """\
## 4. Auto-detect GPU profile

Picks precision + attention backend based on the actual GPU:
- **Turing (T4, V100)** → fp16 + SDPA
- **Ampere+ / Ada (3090, L4, A10, A100, H100)** → bf16 + FA2 (if flash-attn installed; falls back to SDPA otherwise)
"""))


CELLS.append(code("inf-gpu", """\
import torch

assert torch.cuda.is_available(), 'CUDA not available.'
_p       = torch.cuda.get_device_properties(0)
_cap     = (_p.major, _p.minor)
_vram_gb = _p.total_memory / 1e9
_bf16_ok = torch.cuda.is_bf16_supported()
_fa2_ok  = _cap >= (8, 0)

_flash_attn_installed = False
if _fa2_ok:
    try:
        import flash_attn  # noqa: F401
        _flash_attn_installed = True
    except Exception:
        _flash_attn_installed = False

PROFILE = dict(
    label                  = _p.name,
    torch_dtype            = 'bfloat16' if _bf16_ok else 'float16',
    bnb_4bit_compute_dtype = 'bfloat16' if _bf16_ok else 'float16',
    attn_implementation    = 'flash_attention_2' if (_fa2_ok and _flash_attn_installed) else 'sdpa',
)

print(f'GPU         : {_p.name}  ({_vram_gb:.1f} GB)')
print(f'Compute cap : sm_{_cap[0]}{_cap[1]}   bf16 ok: {_bf16_ok}   FA2 wheel: {_flash_attn_installed}')
print(f'-> precision: {PROFILE["torch_dtype"]}   attn: {PROFILE["attn_implementation"]}')
"""))


CELLS.append(md("inf-cfg-md", "## 5. Build configs"))


CELLS.append(code("inf-cfg", """\
from omegaconf import OmegaConf

if SAVED_MODEL_CFG.is_file():
    model_cfg = OmegaConf.load(SAVED_MODEL_CFG)
    print(f'model_cfg <- {SAVED_MODEL_CFG}')
else:
    model_cfg = OmegaConf.load(PROJECT / 'configs' / 'model_config.yaml')
    print(f'model_cfg <- repo default')

if SAVED_TRAIN_CFG.is_file():
    train_cfg = OmegaConf.load(SAVED_TRAIN_CFG)
else:
    train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')

# Compute from auto-detected PROFILE (previous cell).
model_cfg.llm.load_in_4bit              = True
model_cfg.llm.load_in_8bit              = False
model_cfg.llm.attn_implementation       = PROFILE['attn_implementation']
model_cfg.llm.torch_dtype               = PROFILE['torch_dtype']
model_cfg.llm.bnb_4bit_compute_dtype    = PROFILE['bnb_4bit_compute_dtype']
model_cfg.llm.bnb_4bit_quant_type       = 'nf4'
model_cfg.llm.bnb_4bit_use_double_quant = True
model_cfg.llm.gradient_checkpointing    = False    # inference only

# CheXpert classifier — enable iff its checkpoint is pulled.
if CHEXPERT_PT.is_file():
    model_cfg.chexpert_classifier.enabled = True
    print(f'CheXpert classifier checkpoint found -> ENABLED')
else:
    model_cfg.chexpert_classifier.enabled = False
    print('No CheXpert classifier in this run -> DISABLED '
          '(structured_findings will be None unless you pass it explicitly)')

print('\\nmodel_cfg.llm:')
print(OmegaConf.to_yaml(model_cfg.llm))
"""))


CELLS.append(md("inf-load-md", """\
## 6. Load model into kernel

This is the slow phase on cold cache: encoder + 4-bit Vicuna shard download + quantization. Run **once per session** and reuse the `model` object for all subsequent inferences.

Approx timing: ~8 min on T4 cold, ~3 min on L4 cold, seconds when cache is warm.
"""))


CELLS.append(code("inf-load", """\
import time, torch
from model import CXRVisionLanguageModel
from model.rad_dino import BioViLTEncoder
from utils.checkpoint import load_checkpoint

print('[1/3] Building model (Vicuna-7B 4-bit + RAD-DINO + LoRA)…')
print('      On T4 cold cache this takes ~5–8 min for the shard download + quantization.')
_t0 = time.time()
model = CXRVisionLanguageModel(model_cfg)
print(f'      built in {time.time()-_t0:.1f}s', flush=True)

print(f'[2/3] Loading checkpoint from {CKPT_DIR_PULLED} …')
_t0 = time.time()
# Pass the directory, not the .pt file: load_checkpoint splits suffix off
# the path stem, so passing checkpoint_projection.pt makes it look for
# checkpoint_projection_projection.pt (which does not exist) and silently
# skips both projection AND lora — model falls back to raw 4-bit Vicuna.
load_checkpoint(model, str(CKPT_DIR_PULLED))
print(f'      loaded in {time.time()-_t0:.1f}s', flush=True)

print('[3/3] Moving to cuda + eval()')
_t0 = time.time()
model = model.to('cuda').eval()
print(f'      ready in {time.time()-_t0:.1f}s', flush=True)

# Cache transform + chexpert helpers for reuse.
TRANSFORM = BioViLTEncoder.get_transform('val')
print('\\nModel ready. VRAM used:',
      f'{torch.cuda.memory_allocated()/1e9:.2f} GB / {torch.cuda.get_device_properties(0).total_memory/1e9:.1f} GB')
"""))


CELLS.append(md("inf-helpers-md", """\
## 7. Inference helpers

Two convenience functions:

- `predict(image, task, question=None, structured_findings=None, ...)` — single image, single task.
- `predict_report(image, ...)` — cascade: generates findings first, then feeds the model's own findings as context for impression. Returns both sections plus the merged report.

`image` can be a filesystem path (str/Path) or a `PIL.Image.Image`.

### About `structured_findings` (the PNU prompt block)

Your training run baked **oracle CheXpert labels from the CSV/manifest** into the prompt as the `Positive / Negative / Uncertain Abnormalities` block. At inference time you don't have GT labels — you have three choices:

1. **`None` (default)** — no PNU block prepended. The model sees only the image + the task instruction. This is the realistic deployment mode.
2. **Pass a PNU string yourself** — for oracle ablation or testing how much PNU helps:
   ```
   structured_findings='Positive Abnormalities: Cardiomegaly\\nNegative Abnormalities: No Finding, ...\\nUncertain Abnormalities: None'
   ```
3. **Train a CheXpert classifier** (Stage 0 in the training pipeline) and load it via `model_cfg.chexpert_classifier.checkpoint` — your current runs do not have one, so this path is disabled in section 4.
"""))


CELLS.append(code("inf-helpers", """\
import torch
from pathlib import Path
from typing import Optional, Union
from PIL import Image

from data.prompt_templates import (
    build_findings_prompt, build_impression_prompt,
    build_report_prompt,   build_vqa_prompt,
)

ImageLike = Union[str, Path, Image.Image]


def _to_tensor(image: ImageLike) -> torch.Tensor:
    '''Load + transform a CXR into a (1, C, H, W) tensor on cuda.'''
    if isinstance(image, (str, Path)):
        img = Image.open(image).convert('RGB')
    else:
        img = image.convert('RGB') if image.mode != 'RGB' else image
    t = TRANSFORM(img)                                # (C, H, W)
    return t.unsqueeze(0).to('cuda')                  # (1, C, H, W)


@torch.no_grad()
def predict(
    image:               ImageLike,
    task:                str,
    question:            Optional[str] = None,
    structured_findings: Optional[str] = None,
    max_new_tokens:      int = None,
    temperature:         float = None,
    do_sample:           bool = None,
    num_beams:           int = None,
) -> str:
    '''Run the model on one image.

    task ∈ {'findings', 'impression', 'report', 'vqa'}.
    For task='vqa', `question` is required.
    '''
    assert task in ('findings', 'impression', 'report', 'vqa'), f'bad task: {task}'
    if task == 'vqa':
        assert question, 'VQA requires `question`.'

    sf = structured_findings
    if task == 'findings':
        prompt = build_findings_prompt(sf, randomize=False)
    elif task == 'impression':
        prompt = build_impression_prompt(sf, randomize=False)
    elif task == 'report':
        prompt = build_report_prompt(sf, randomize=False)
    else:
        prompt = build_vqa_prompt(question, sf)

    out = model.generate(
        images         = _to_tensor(image),
        prompts        = [prompt],
        max_new_tokens = max_new_tokens if max_new_tokens is not None else MAX_NEW_TOKENS,
        temperature    = temperature    if temperature    is not None else TEMPERATURE,
        do_sample      = do_sample      if do_sample      is not None else DO_SAMPLE,
        num_beams      = num_beams      if num_beams      is not None else NUM_BEAMS,
    )
    return out[0]


@torch.no_grad()
def predict_report(
    image:               ImageLike,
    structured_findings: Optional[str] = None,
    max_new_tokens:      int = None,
):
    '''Cascade: generate findings first, then feed them as context for impression.

    Matches the `split_cascade` training recipe but with model-generated (not
    GT) findings — i.e. a real end-to-end cascade (no GT leakage).

    Returns dict {findings, impression, report}.
    '''
    # 1) Findings (uses PNU context only — None by default)
    findings = predict(image, task='findings',
                       structured_findings=structured_findings,
                       max_new_tokens=max_new_tokens)

    # 2) Impression: feed the generated findings as the structured_findings
    #    context — same shape as the cascade trainer saw at training time
    #    but with the model's own findings instead of GT.
    impression = predict(image, task='impression',
                         structured_findings=findings,
                         max_new_tokens=max_new_tokens)

    return {
        'findings':   findings,
        'impression': impression,
        'report':     f'Findings: {findings}\\n\\nImpression: {impression}',
    }
"""))


CELLS.append(md("inf-show-md", "## 8. Display helper (optional, pretty-prints image + outputs)"))


CELLS.append(code("inf-show", """\
import matplotlib.pyplot as plt
from PIL import Image as _PILImage

def show(image, title=None, max_size=512):
    '''Render an image inline with optional title.'''
    if isinstance(image, (str, Path)):
        img = _PILImage.open(image).convert('RGB')
    else:
        img = image
    img.thumbnail((max_size, max_size))
    plt.figure(figsize=(6, 6))
    plt.imshow(img)
    plt.axis('off')
    if title:
        plt.title(title)
    plt.show()

def pretty(d: dict, header: str = None):
    '''Pretty-print a dict of {label: text}.'''
    if header:
        print('=' * 80)
        print(header)
    for k, v in d.items():
        print('-' * 80)
        print(f'[{k.upper()}]')
        print(v)
    print('=' * 80)
"""))


CELLS.append(md("inf-ex1-md", """\
## 9. Examples

### Example A — single image, all tasks

Edit `IMAGE_PATH` and `VQA_QUESTION` below, then run.
"""))


CELLS.append(code("inf-ex1", """\
# Pick an image — anywhere on disk. Colab: drag-and-drop into the Files pane
# (left sidebar) and use '/content/<filename>'. Local: any path that exists.
IMAGE_PATH = '/content/sample_cxr.jpg'

# Optional PNU override (see helpers section for format). None → no PNU block.
STRUCTURED_FINDINGS = None

# A clinical question for the VQA call below.
VQA_QUESTION = 'Is there any pleural effusion visible in this chest X-ray?'

if not Path(IMAGE_PATH).is_file():
    print(f'!! {IMAGE_PATH} does not exist. Upload an image or change IMAGE_PATH.')
else:
    show(IMAGE_PATH, title=Path(IMAGE_PATH).name)

    if STRUCTURED_FINDINGS:
        print('Structured findings (PNU) override:')
        print(STRUCTURED_FINDINGS)
        print()

    findings   = predict(IMAGE_PATH, task='findings',
                         structured_findings=STRUCTURED_FINDINGS)
    impression = predict(IMAGE_PATH, task='impression',
                         structured_findings=STRUCTURED_FINDINGS)
    vqa_ans    = predict(IMAGE_PATH, task='vqa', question=VQA_QUESTION,
                         structured_findings=STRUCTURED_FINDINGS)

    pretty({
        'findings':       findings,
        'impression':     impression,
        f'vqa ({VQA_QUESTION!r})': vqa_ans,
    }, header=Path(IMAGE_PATH).name)
"""))


CELLS.append(md("inf-ex2-md", """\
### Example B — cascade report (findings → impression on model output)
"""))


CELLS.append(code("inf-ex2", """\
if Path(IMAGE_PATH).is_file():
    out = predict_report(IMAGE_PATH)
    show(IMAGE_PATH, title='Cascade report')
    pretty(out, header=f'Cascade: {Path(IMAGE_PATH).name}')
"""))


CELLS.append(md("inf-ex3-md", """\
### Example C — batch a folder of images

Drops all results into `BATCH_OUT_JSON` for downstream analysis.
"""))


CELLS.append(code("inf-ex3", """\
import json
from tqdm.auto import tqdm

BATCH_FOLDER   = '/content/cxr_inbox'            # folder of *.jpg / *.png
BATCH_TASK     = 'report'                        # 'findings' | 'impression' | 'report' | 'vqa' | 'cascade'
BATCH_QUESTION = None                            # only used when BATCH_TASK == 'vqa'
BATCH_OUT_JSON = WORK / f'inference_{RUN_ID}_{BATCH_TASK}.json'

EXTS = {'.jpg', '.jpeg', '.png'}
folder = Path(BATCH_FOLDER)
imgs = sorted(p for p in folder.rglob('*') if p.suffix.lower() in EXTS) \\
       if folder.is_dir() else []
print(f'{len(imgs)} image(s) under {folder}')

results = []
for p in tqdm(imgs, desc=f'Inference [{BATCH_TASK}]'):
    try:
        if BATCH_TASK == 'cascade':
            out = predict_report(p)
        else:
            out = {BATCH_TASK: predict(p, task=BATCH_TASK,
                                       question=BATCH_QUESTION)}
        results.append({'image': str(p), **out})
    except Exception as e:
        results.append({'image': str(p), 'error': f'{type(e).__name__}: {e}'})

if results:
    BATCH_OUT_JSON.write_text(json.dumps(results, indent=2))
    print(f'Wrote {len(results)} entries -> {BATCH_OUT_JSON}')
"""))


CELLS.append(md("inf-done-md", """\
### Tips

- Loading the model takes 5–8 min on T4 cold cache. After that, each generation is ~3–8 s.
- To switch to a different trained run: change `RUN_ID` in section 0, re-run sections 3–5 (skip the env / pip sections).
- For deterministic outputs keep `DO_SAMPLE=False`. For diversity, set `DO_SAMPLE=True, TEMPERATURE=0.7`.
- `predict_report()` is the realistic end-to-end pipeline (no GT leakage). `predict(task='report')` only works on runs trained with `report_mode='merged'`.
"""))


# ─────────────────────────────────────────────────────────────────────

nb = {
    "cells": CELLS,
    "metadata": {
        "accelerator": "GPU",
        "colab": {
            "gpuType": "T4",
            "provenance": [],
            "machine_shape": "hm",
        },
        "kernelspec": {
            "display_name": "Python 3",
            "name": "python3",
        },
        "language_info": {
            "name": "python",
            "version": "3.10",
        },
    },
    "nbformat": 4,
    "nbformat_minor": 5,
}

NB_PATH.write_text(json.dumps(nb, indent=1, ensure_ascii=False), encoding="utf-8")
print(f"wrote {NB_PATH}  ({len(CELLS)} cells)")