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

Run:
    python scripts/_build_eval_notebook.py
"""

import json
from pathlib import Path

NB_PATH = Path(__file__).parent / "cxrvlm_colab_eval.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),
    }


# ─────────────────────────────────────────────────────────────────────
# Cell sources
# ─────────────────────────────────────────────────────────────────────

CELLS = []


CELLS.append(md("eval-0", """\
# CXR-VLM — Evaluation Notebook (Colab T4)

Standalone evaluation + inference for a trained CXR-VLM run.

What it does:
1. Pulls project code from `<HF_USER>/cxr-vlm-code`.
2. Pulls the chosen dataset from `<HF_USER>/cxr-vlm-data` (same layout the trainer used).
3. Pulls a trained run checkpoint from `<HF_USER>/cxr-vlm-runs/<RUN_ID>/stage2/{best|last}/`.
4. Pulls the run's config snapshot (`<RUN_ID>/configs/{model,train}_config.yaml`) so the model is rebuilt **with the exact same architecture / report_mode / image_mode** as training.
5. **Auto-detects the GPU** (T4 / L4 / 3090 / A10 / A100 / H100) and patches the configs accordingly — fp16+SDPA on Turing, bf16+FA2 on Ampere/Ada, batch size scaled to VRAM. Also patches local dataset paths.
6. Runs `python -m evaluation.evaluate` on the **test split** for all available tasks.
7. Saves predictions + metrics under `RESULTS_DIR/<RUN_ID>/` (and optionally uploads them to `<HF_USER>/cxr-vlm-runs/<RUN_ID>/results/`).

Set the variables in **section 0** and run all cells top-to-bottom.

**Want interactive inference on individual images?** Use the companion notebook **`cxrvlm_colab_inference.ipynb`** — same model pull, but with image preview + free-form prompts.
"""))


CELLS.append(md("eval-select-md", """\
## 0. Select run + dataset + options

Change the variables in the cell below. `RUN_ID` decides which trained model is pulled; everything else (dataset name, report/image mode) defaults to whatever was used at training time (read from the run's config snapshot on HF), but you can override.
"""))


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

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

# ── Which trained run to evaluate ────────────────────────────────
# This MUST be an existing folder on HF_RUNS_REPO.
# Example: 'MIMIC-CXR_resized_run_1' | 'IU-Xray_run_2' | 'MIMIC-CXR_run_3'
RUN_ID        = 'MIMIC-CXR_resized_run_1'

# Which stage-2 checkpoint to load.
#   'best' → {RUN_ID}/stage2/best/  (final / best eval_loss)
#   'last' → {RUN_ID}/stage2/last/  (most recent intermediate save)
CKPT_PICK     = 'best'

# ── Dataset (auto-derived from RUN_ID prefix; override if needed) ──
# Supported: 'MIMIC-CXR' | 'MIMIC-CXR_resized' | 'IU-Xray'
DATASET_NAME  = None        # None → auto-detect from RUN_ID prefix
REPORT_MODE   = None        # None → read from run's saved train_config.yaml
IMAGE_MODE    = None        # None → read from run's saved train_config.yaml

# ── What to evaluate ─────────────────────────────────────────────
TASK          = 'all'       # 'all' | 'findings' | 'impression' | 'report' | 'vqa'
SPLIT         = 'test'      # always 'test' for this notebook; here for visibility
BATCH_SIZE    = None        # None → auto from GPU profile (T4:1, L4:4, A100:8, H100:16)
MAX_NEW_TOKENS = 300

# ── Metric config (cross-paper comparability) ────────────────────
# BERTScore: roberta-large + rescale → low, paper-comparable scores (~0.3-0.5).
#   Set BERTSCORE_RESCALE=False (and/or 'distilbert-base-uncased') for the old
#   raw scores (~0.8) — those are NOT comparable across papers.
# METEOR: 'nltk' is easy but scores higher than papers; 'pycoco' is comparable
#   (needs Java + `pip install pycocoevalcap`).
BERTSCORE_MODEL   = 'roberta-large'
BERTSCORE_RESCALE = True
METEOR_IMPL       = 'nltk'   # 'nltk' | 'pycoco'

# ── LLM-as-judge for VQA (optional; needs OPENAI_API_KEY) ────────
LLM_JUDGE             = False
LLM_JUDGE_MODEL       = 'gpt-4o-mini'
LLM_JUDGE_BASE_URL    = None          # e.g. Gemini OpenAI-compat endpoint
LLM_JUDGE_MAX_SAMPLES = None          # cap cost; None → all VQA samples

# ── Output ───────────────────────────────────────────────────────
# Local folder where predictions_*.json + metrics_summary.json land.
# Files end up at  {LOCAL_RESULTS_DIR}/{RUN_ID}/...
LOCAL_RESULTS_DIR = 'results'

# Push the results folder back to HF_RUNS_REPO under
#   {RUN_ID}/results/predictions_*.json + metrics_summary.json
UPLOAD_RESULTS_TO_HF = True

# ── Auto-derive DATASET_NAME from RUN_ID prefix if not set ───────
if DATASET_NAME is None:
    for cand in ('MIMIC-CXR_resized', 'MIMIC-CXR', 'IU-Xray'):
        if RUN_ID.startswith(cand + '_run_'):
            DATASET_NAME = cand
            break
    assert DATASET_NAME is not None, \\
        f"Cannot auto-derive DATASET_NAME from RUN_ID={RUN_ID!r}. " \\
        "Expected prefix one of: MIMIC-CXR_resized / MIMIC-CXR / IU-Xray. Set DATASET_NAME explicitly."

assert PLATFORM     in ('kaggle', 'colab', 'lightning', 'gcp', 'local')
assert DATASET_NAME in ('MIMIC-CXR', 'MIMIC-CXR_resized', 'IU-Xray')
assert CKPT_PICK    in ('best', 'last')
assert TASK         in ('all', 'findings', 'impression', 'report', 'vqa')

print(f'PLATFORM      = {PLATFORM}')
print(f'RUN_ID        = {RUN_ID}     (ckpt: stage2/{CKPT_PICK})')
print(f'DATASET_NAME  = {DATASET_NAME}')
print(f'TASK          = {TASK}   SPLIT = {SPLIT}')
print(f'LOCAL_RESULTS_DIR = {LOCAL_RESULTS_DIR}   (upload→HF: {UPLOAD_RESULTS_TO_HF})')
"""))


CELLS.append(md("eval-env-md", """\
## 1. Environment + pip install (matches the training notebook)
"""))


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

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


CELLS.append(code("eval-pip", """\
import os as _os

# Same dependency set as the training notebook — keeps load_checkpoint /
# evaluate.py running against the exact stack the model was trained with.
!pip uninstall -y -q torchao transformers bitsandbytes peft accelerate

# Let pip pick latest bnb that matches Colab's CUDA + triton.
!pip install -q -U bitsandbytes

# Pin transformers / peft to the same window the trainer uses.
!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' \\
    nltk rouge-score bert-score sacrebleu

# 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.
# Skipped entirely on Turing (T4/V100) since FA2 requires sm_80+.
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("eval-versions", """\
import torch, transformers, bitsandbytes, peft, accelerate, huggingface_hub, httpx

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

# httpx 0.28+ compat shim — same patch as the training notebook.
# transformers <=4.50 calls httpx.Client.head(..., allow_redirects=True)
# which httpx 0.28 removed; translate the kwarg at the call site.
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
    print(f'httpx {httpx.__version__}: monkey-patched allow_redirects -> follow_redirects')

_patch_httpx()

assert torch.cuda.is_available(), 'CUDA not available — refusing to evaluate on CPU.'
_p = torch.cuda.get_device_properties(0)
print(f'\\nGPU         : {_p.name}  ({_p.total_memory/1e9:.1f} GB)')
print(f'Compute cap : sm_{_p.major}{_p.minor}  (BF16 ok: {torch.cuda.is_bf16_supported()})')
"""))


CELLS.append(md("eval-paths-md", """\
## 2. Paths + pull code + pull dataset

Identical to the training notebook for sections that overlap (HF code + per-dataset payload).
"""))


CELLS.append(code("eval-paths", """\
# ── 1) 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':
    WORK = Path('/workspace')
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, hf_hub_download, HfApi

# ── 2) Code: flat folder, 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'),
))

# ── 3) Data: layout depends on DATASET_NAME (same logic as train notebook) ──
DATA_SRC = WORK / 'data'
DATA_SRC.mkdir(parents=True, exist_ok=True)

if DATASET_NAME == 'MIMIC-CXR_resized':
    import tarfile
    mr_dir = DATA_SRC / 'MIMIC-CXR_resized'
    mr_dir.mkdir(parents=True, exist_ok=True)
    files_dir = mr_dir / 'files'

    manifests_present = all(
        (mr_dir / f).is_file() for f in ('manifest_train.csv', 'manifest_val.csv', 'manifest_test.csv')
    )
    if manifests_present and files_dir.is_dir() and any(files_dir.glob('p*')):
        print(f'{mr_dir} already populated — skipping download.')
    else:
        api = HfApi(token=os.environ['HF_TOKEN'])
        all_files = api.list_repo_files(repo_id=HF_DATA_REPO, repo_type='dataset')
        mr_files  = [f for f in all_files if f.startswith('MIMIC-CXR_resized/')]
        tar_files = sorted(f for f in mr_files if f.endswith('.tar'))
        meta_files = [f for f in mr_files if not f.endswith('.tar')]
        print(f'MIMIC-CXR_resized on HF: {len(tar_files)} tar shards + {len(meta_files)} metadata files')

        # Manifests / vqa / SHARDS.txt
        snapshot_download(
            repo_id        = HF_DATA_REPO,
            repo_type      = 'dataset',
            allow_patterns = ['MIMIC-CXR_resized/*.csv',
                              'MIMIC-CXR_resized/*.json',
                              'MIMIC-CXR_resized/*.txt',
                              'MIMIC-CXR_resized/vqa/**'],
            token          = os.environ['HF_TOKEN'],
            local_dir      = str(DATA_SRC),
        )

        # Tar shards — sequential extract + delete to keep peak disk low.
        for i, tf in enumerate(tar_files, 1):
            print(f'  [{i}/{len(tar_files)}] {tf}')
            tar_path = Path(hf_hub_download(
                repo_id=HF_DATA_REPO, repo_type='dataset',
                filename=tf, token=os.environ['HF_TOKEN'],
                local_dir=str(DATA_SRC),
            ))
            with tarfile.open(tar_path) as t:
                t.extractall(mr_dir)
            tar_path.unlink(missing_ok=True)
        print(f'  done. {mr_dir} ready.')

else:
    # MIMIC-CXR / IU-Xray — single zip per dataset.
    import zipfile
    zip_name = f'{DATASET_NAME}.zip'
    marker   = DATA_SRC / DATASET_NAME

    if not marker.exists():
        print(f'Pulling {zip_name} from HF …')
        zpath = hf_hub_download(
            repo_id   = HF_DATA_REPO,
            filename  = zip_name,
            repo_type = 'dataset',
            token     = os.environ['HF_TOKEN'],
            local_dir = str(DATA_SRC),
        )
        print(f'  unzipping -> {DATA_SRC}')
        with zipfile.ZipFile(zpath) as zf:
            zf.extractall(DATA_SRC)
        try:
            os.remove(zpath)
        except OSError:
            pass
    else:
        print(f'{marker} already present — skipping download.')

print(f'Contents of {DATA_SRC}: {sorted(os.listdir(DATA_SRC))}')

# ── 4) Copy code into writable PROJECT dir + chdir ─────────────────
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('CODE_SRC :', CODE_SRC)
print('DATA_SRC :', DATA_SRC)
print('PROJECT  :', PROJECT)
print('WORK     :', WORK)
"""))


CELLS.append(md("eval-find-md", """\
## 3. Locate the dataset payload on disk
"""))


CELLS.append(code("eval-find", """\
# Same dataset-payload finders as the training notebook. Filling in the
# variables the config cell below will consume.

def find_split_parent(root: Path) -> Path:
    for cand in [root, root / 'MIMIC-CXR', root / 'data' / 'MIMIC-CXR']:
        if (cand / 'train').exists() and (cand / 'valid').exists() and (cand / 'test').exists():
            return cand
    for p in root.rglob('train'):
        if p.is_dir() and (p.parent / 'valid').exists() and (p.parent / 'test').exists():
            return p.parent
    raise FileNotFoundError('Could not find train/ valid/ test/ under ' + str(root))


def find_mimic_resized_root(root: Path) -> Path:
    for cand in [root / 'MIMIC-CXR_resized', root, *root.rglob('MIMIC-CXR_resized')]:
        if (cand / 'manifest_train.csv').is_file():
            return cand
    raise FileNotFoundError(
        f'Could not find MIMIC-CXR_resized payload under {root}. '
        'Expected manifest_train.csv (alongside manifest_val.csv / manifest_test.csv).'
    )


def find_iu_dirs(root: Path):
    for cand in [root / 'IU-Xray', *root.rglob('IU-Xray')]:
        if not cand.is_dir():
            continue
        imgs = cand / 'images'
        lbls = cand / 'labels'
        if imgs.is_dir() and lbls.is_dir() and any(lbls.glob('*.xml')):
            return imgs, lbls
        legacy = lbls / 'ecgen-radiology'
        if imgs.is_dir() and legacy.is_dir() and any(legacy.glob('*.xml')):
            return imgs, legacy
    img_dir = lbl_dir = None
    for cand in [root / 'images', *root.rglob('images')]:
        if cand.is_dir() and any(cand.glob('CXR*.png')):
            img_dir = cand; break
    for cand in [root / 'labels', *root.rglob('labels')]:
        if cand.is_dir() and any(cand.glob('*.xml')):
            lbl_dir = cand; break
    if lbl_dir is None:
        for cand in root.rglob('ecgen-radiology'):
            if cand.is_dir() and any(cand.glob('*.xml')):
                lbl_dir = cand; break
    return img_dir, lbl_dir


CXR_ROOT      = None
VQA_ROOT      = None
MR_ROOT       = None
IU_IMAGES_DIR = None
IU_LABELS_DIR = None

if DATASET_NAME == 'MIMIC-CXR':
    CXR_ROOT = find_split_parent(DATA_SRC)
    print('MIMIC-CXR root:', CXR_ROOT)
    for s in ('train', 'valid', 'test'):
        d = CXR_ROOT / s
        assert d.exists(), f'Missing split dir: {d}'
        print(f'  {s:<6s} -> {d}')

    for p in DATA_SRC.rglob('MIMIC-Ext-MIMIC-CXR-VQA'):
        cand = p / 'dataset'
        if cand.exists() and (cand / 'train.json').exists():
            VQA_ROOT = cand; break
    if VQA_ROOT is None:
        print('VQA root: NOT FOUND -> VQA task will be skipped')
    else:
        print('VQA root:', VQA_ROOT)

elif DATASET_NAME == 'MIMIC-CXR_resized':
    MR_ROOT = find_mimic_resized_root(DATA_SRC)
    print('MIMIC-CXR_resized root:', MR_ROOT)
    for cf in ('manifest_train.csv', 'manifest_val.csv', 'manifest_test.csv'):
        f = MR_ROOT / cf
        print(f'  {cf}: {"OK" if f.is_file() else "MISSING"}')
    for sub in ('files', 'vqa'):
        d = MR_ROOT / sub
        print(f'  {sub:<5s}: {"OK" if d.is_dir() else "MISSING"}  ({d})')

else:  # IU-Xray
    IU_IMAGES_DIR, IU_LABELS_DIR = find_iu_dirs(DATA_SRC)
    assert IU_IMAGES_DIR is not None, f'IU images/ not found under {DATA_SRC}'
    assert IU_LABELS_DIR is not None, f'IU labels/ (with *.xml) not found under {DATA_SRC}'
    print('IU images dir:', IU_IMAGES_DIR, '->', len(list(IU_IMAGES_DIR.glob('*.png'))), 'PNGs')
    print('IU labels dir:', IU_LABELS_DIR, '->', len(list(IU_LABELS_DIR.glob('*.xml'))), 'XMLs')
"""))


CELLS.append(md("eval-pull-run-md", """\
## 4. Pull trained checkpoint + run config snapshot from HF Runs repo

Layout on `HF_RUNS_REPO`:

```
{RUN_ID}/
  configs/{model,train}_config.yaml   ← snapshot taken at training time
  stage2/best/   checkpoint_projection.pt + checkpoint_lora/ [+ checkpoint_chexpert_classifier.pt]
  stage2/last/   same shape, intermediate save
```

We pull `{RUN_ID}/configs/` and `{RUN_ID}/stage2/{CKPT_PICK}/` into `{WORK}/run_pull/{RUN_ID}/`.
"""))


CELLS.append(code("eval-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
assert RUN_DIR_PULLED.is_dir(), f'Pull failed — {RUN_DIR_PULLED} missing.'

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}. ' \\
    f'Check that {RUN_ID}/stage2/{CKPT_PICK}/ exists on {HF_RUNS_REPO}.'
assert (LORA_DIR / 'adapter_config.json').is_file(), \\
    f'LoRA adapter_config.json not found in {LORA_DIR}. Stage-2 checkpoint partial?'

print()
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 configs (may or may not exist on older runs)
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'
print()
print(f'  saved train_cfg : {SAVED_TRAIN_CFG}  (exists: {SAVED_TRAIN_CFG.is_file()})')
print(f'  saved model_cfg : {SAVED_MODEL_CFG}  (exists: {SAVED_MODEL_CFG.is_file()})')
"""))


CELLS.append(md("eval-gpu-md", """\
## 5. Auto-detect GPU profile

Mirrors the training notebook: picks precision (bf16/fp16), attention backend (FA2/SDPA), and an eval batch size based on the actual GPU's compute capability + VRAM.

| Bucket | Examples | Precision | Attn | Eval batch |
|---|---|---|---|---|
| 70+ GB | A100/H100 80GB | bf16 | FA2 (if installed) | 16 |
| 35–69 GB | A100 40GB | bf16 | FA2 | 8 |
| 22–34 GB | 3090 / L4 / A10 24GB | bf16 | FA2 | 4 |
| 14–21 GB | T4 / V100 16GB | fp16 | SDPA | 1 |

The eval batch sizes are smaller than training's because generation builds a KV cache that scales with `batch × max_new_tokens`. Override by setting `BATCH_SIZE` in section 0 to a concrete number.
"""))


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

assert torch.cuda.is_available(), 'CUDA not available — refusing to write a CPU profile.'
_props   = torch.cuda.get_device_properties(0)
_cap     = (_props.major, _props.minor)
_vram_gb = _props.total_memory / 1e9
_bf16_ok = torch.cuda.is_bf16_supported()
_fa2_ok  = _cap >= (8, 0)        # FA2 needs Ampere+ (sm_80 or newer)

# Detect whether flash-attn package is actually importable. FA2 falls back to
# SDPA inside cxr_vlm.py if missing, but knowing here lets us print clearly.
_flash_attn_installed = False
if _fa2_ok:
    try:
        import flash_attn  # noqa: F401
        _flash_attn_installed = True
    except Exception:
        _flash_attn_installed = False

print(f'GPU         : {_props.name}  ({_vram_gb:.1f} GB)')
print(f'Compute cap : sm_{_cap[0]}{_cap[1]}')
print(f'BF16 native : {_bf16_ok}')
print(f'FA2 capable : {_fa2_ok}   flash-attn installed: {_flash_attn_installed}')

# Eval batch size — smaller than training because generation builds a KV cache
# that scales with batch × max_new_tokens.
if _vram_gb >= 70:
    GPU_LABEL, _EVAL_BS, _NW = 'A100/H100 80GB',         16, 16
elif _vram_gb >= 35:
    GPU_LABEL, _EVAL_BS, _NW = 'A100 40GB',               8, 12
elif _vram_gb >= 22:
    GPU_LABEL, _EVAL_BS, _NW = 'RTX 3090 / L4 / A10 24GB', 4, 8
elif _vram_gb >= 14:
    GPU_LABEL, _EVAL_BS, _NW = 'T4 / V100 (15-16GB)',     1, 2
else:
    GPU_LABEL, _EVAL_BS, _NW = f'unknown ({_vram_gb:.0f}GB) - conservative', 1, 2

PROFILE = dict(
    label                       = GPU_LABEL,
    bf16                        = bool(_bf16_ok),
    fp16                        = not _bf16_ok,
    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',
    per_device_eval_batch_size  = _EVAL_BS,
    dataloader_num_workers      = _NW,
)

# Allow the section-0 selector to override (e.g. BATCH_SIZE=8 to push harder).
if BATCH_SIZE is not None:
    PROFILE['per_device_eval_batch_size'] = int(BATCH_SIZE)
    print(f'BATCH_SIZE override: section 0 forced batch={BATCH_SIZE}')

# Final BATCH_SIZE the rest of the notebook (eval-run cell) will use.
BATCH_SIZE = PROFILE['per_device_eval_batch_size']

print(f'\\n-> Profile  : {PROFILE["label"]}')
for k, v in PROFILE.items():
    if k == 'label': continue
    print(f'    {k:<32s} = {v}')
"""))


CELLS.append(md("eval-cfg-md", """\
## 6. Build configs

Strategy: start from the run's saved config snapshot if present (so `report_mode`, `image_mode`, `lora.r`, `num_image_tokens`, etc. match training). If absent, fall back to the repo defaults. Then patch:
- dataset paths to the local download
- compute (precision + attn backend + batch) from the auto-detected `PROFILE`
- HF Hub tracker → uses `HF_RUNS_REPO` so the results upload lands under `{RUN_ID}/results/`
- pin `run_id` to `RUN_ID` so `evaluate.py` writes under `{LOCAL_RESULTS_DIR}/{RUN_ID}/`
"""))


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

# ── 1) Load base configs: prefer the run's saved snapshot ─────────
if SAVED_TRAIN_CFG.is_file():
    train_cfg = OmegaConf.load(SAVED_TRAIN_CFG)
    print(f'train_cfg <- {SAVED_TRAIN_CFG}')
else:
    train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')
    print(f'train_cfg <- repo default (no snapshot on HF)')

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 (no snapshot on HF)')

# ── 2) Allow notebook overrides for report/image mode ─────────────
if REPORT_MODE is not None:
    train_cfg.data.report_mode = REPORT_MODE
if IMAGE_MODE is not None:
    train_cfg.data.image_mode = IMAGE_MODE
print(f'report_mode = {train_cfg.data.report_mode}   image_mode = {train_cfg.data.image_mode}')

# ── 3) Patch dataset paths for the local download ────────────────
train_cfg.data.dataset_name = DATASET_NAME

if DATASET_NAME == 'MIMIC-CXR':
    train_cfg.data.mimic_cxr_root   = str(CXR_ROOT)
    train_cfg.data.mimic_auto_build = True

    _cx = (sorted(DATA_SRC.rglob('*chexpert*.csv'))
           or sorted(DATA_SRC.rglob('*chexbert*.csv')))
    train_cfg.data.mimic_chexpert_csv = str(_cx[0]) if _cx else None
    train_cfg.data.mimic_vqa_root     = str(VQA_ROOT) if VQA_ROOT is not None else None

    out_dir = PROJECT / 'data' / 'data_files'
    out_dir.mkdir(parents=True, exist_ok=True)
    train_cfg.data.instruct_json = str(out_dir / 'mimic_cxr_instruct_unified.json')

elif DATASET_NAME == 'MIMIC-CXR_resized':
    train_cfg.data.mimic_cxr_resized.root          = str(MR_ROOT)
    train_cfg.data.mimic_cxr_resized.manifest_dir  = None
    train_cfg.data.mimic_cxr_resized.vqa_dir       = None
    train_cfg.data.mimic_cxr_resized.reports_root  = None
    train_cfg.data.mimic_cxr_resized.auto_build    = True
    out_dir = PROJECT / 'data' / 'data_files'
    out_dir.mkdir(parents=True, exist_ok=True)
    train_cfg.data.mimic_cxr_resized.instruct_json = str(out_dir / 'mimic_cxr_resized_instruct.json')

else:  # IU-Xray
    train_cfg.data.iu_xray.images_dir = str(IU_IMAGES_DIR)
    train_cfg.data.iu_xray.labels_dir = str(IU_LABELS_DIR)
    train_cfg.data.iu_xray.auto_build = True
    out_dir = PROJECT / 'data' / 'data_files'
    out_dir.mkdir(parents=True, exist_ok=True)
    train_cfg.data.iu_xray.instruct_json = str(out_dir / 'iu_xray_instruct.json')

train_cfg.data.train_split = 'train'
train_cfg.data.val_split   = 'validate'
train_cfg.data.test_split  = 'test'

# ── 4) Apply auto-detected compute profile (overrides saved config) ──
train_cfg.training.fp16 = PROFILE['fp16']
train_cfg.training.bf16 = PROFILE['bf16']
train_cfg.training.per_device_train_batch_size = PROFILE['per_device_eval_batch_size']
train_cfg.training.per_device_eval_batch_size  = PROFILE['per_device_eval_batch_size']
train_cfg.training.dataloader_num_workers      = PROFILE['dataloader_num_workers']
train_cfg.training.dataloader_pin_memory       = True

# Disable feature cache for eval (test images haven't been precomputed).
train_cfg.data.feature_cache_dir = None

# 4-bit QLoRA — must match how the trainer set it up for the saved LoRA to load.
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     # eval-only - no backward pass

# ── 5) CheXpert classifier ───────────────────────────────────────
# The training notebook keeps it disabled (oracle PNU from CSV/manifest).
# Only enable here if the trained run actually has a learned classifier
# checkpoint saved alongside.
if CHEXPERT_PT.is_file():
    model_cfg.chexpert_classifier.enabled = True
    print(f'CheXpert classifier checkpoint found at {CHEXPERT_PT} -> enabled')
else:
    model_cfg.chexpert_classifier.enabled = False
    print('No CheXpert classifier checkpoint -> disabled (oracle PNU from CSV/manifest)')

# ── 6) HF Hub tracker — points results uploads at HF_RUNS_REPO ──
CKPT_ROOT = WORK / 'ckpt_eval'
CKPT_ROOT.mkdir(parents=True, exist_ok=True)
train_cfg.training.output_root = str(CKPT_ROOT)

if UPLOAD_RESULTS_TO_HF:
    train_cfg.hf_hub.enabled        = True
    train_cfg.hf_hub.repo_id        = HF_RUNS_REPO
    train_cfg.hf_hub.token_env      = 'HF_TOKEN'
    train_cfg.hf_hub.private        = True
else:
    train_cfg.hf_hub.enabled = False

train_cfg.hf_hub.run_state_file = str(CKPT_ROOT / 'run_id.txt')

# Pin RUN_ID so resolve_run_id picks it up exactly.
Path(train_cfg.hf_hub.run_state_file).write_text(RUN_ID)

# ── 7) Save patched configs into the project so the subprocess sees them ──
OmegaConf.save(train_cfg, PROJECT / 'configs' / 'train_config.yaml')
OmegaConf.save(model_cfg, PROJECT / 'configs' / 'model_config.yaml')

print('--- train_cfg.data ---');    print(OmegaConf.to_yaml(train_cfg.data))
print('--- train_cfg.training ---');print(OmegaConf.to_yaml(train_cfg.training))
print('--- train_cfg.hf_hub ---');  print(OmegaConf.to_yaml(train_cfg.hf_hub))
print('--- model_cfg.llm ---');     print(OmegaConf.to_yaml(model_cfg.llm))
"""))


CELLS.append(md("eval-verify-md", """\
## 7. Verify dataset before running eval

Quick pre-flight check: triggers the instruct-JSON builder (if not cached), then prints **per-split × per-task** sample counts. Catches issues like "VQA = 0 samples" (path-format mismatch in the builder) **before** spending 2h on an eval that has nothing to evaluate.

If `vqa` column shows 0 in the test split:
- the model was likely **not trained on VQA** either (same JSON cache used both ways)
- options: skip VQA via `TASK='findings'` then run a second job with `TASK='impression'`, OR fix the builder + retrain.
"""))


CELLS.append(code("eval-verify", """\
import json as _json
from collections import Counter
from utils.dataset_resolver import resolve_dataset_spec

# Reload the patched config snapshot the eval subprocess will see.
train_cfg = OmegaConf.load(PROJECT / 'configs' / 'train_config.yaml')
spec = resolve_dataset_spec(train_cfg)
print(f'Dataset      : {spec.dataset_name}')
print(f'Instruct JSON: {spec.instruct_json}')
print(f'Image root   : {spec.image_root}')
print(f'Tasks (cfg)  : {spec.tasks}')
print()

# Load the JSON the dataset module will read.
samples = _json.loads(open(spec.instruct_json, encoding='utf-8').read())
print(f'Total samples in JSON: {len(samples):,}')

# Cross-tab: (split, task) -> count
ctab = Counter((s['split'], s['task']) for s in samples)
splits = sorted({k[0] for k in ctab})
tasks  = sorted({k[1] for k in ctab})

# Pretty table
col_w = max(10, max(len(t) for t in tasks) + 2)
hdr = f'{\"split\":<10} | ' + ' | '.join(f'{t:>{col_w}}' for t in tasks) + ' |     total'
print(hdr); print('-' * len(hdr))
for sp in splits:
    row_vals = [ctab.get((sp, t), 0) for t in tasks]
    total = sum(row_vals)
    print(f'{sp:<10} | ' + ' | '.join(f'{v:>{col_w},}' for v in row_vals)
          + f' | {total:>9,}')

# Loud warning if VQA is expected but missing.
test_vqa = ctab.get(('test', 'vqa'), 0)
if 'vqa' in spec.tasks and test_vqa == 0:
    print()
    print('!! WARNING: vqa task is configured but TEST split has 0 vqa samples.')
    print('   This usually means the dataset builder dropped all VQA rows due to')
    print('   path-format mismatch between vqa/*.json and manifest_*.csv.')
    print('   Check the builder log above for the line:')
    print('     [mimic_cxr_resized_builder] vqa added/dropped : N / M')
    print('   If N=0 the model was likely NOT trained on VQA either — same JSON cache.')
elif 'vqa' in spec.tasks:
    print(f'\\nVQA in test split: {test_vqa:,} samples — OK')
"""))


CELLS.append(md("eval-run-md", """\
## 8. Run evaluation

Calls `python -m evaluation.evaluate` as a subprocess. It will:
1. Build the unified instruct JSON for the chosen `(report_mode, image_mode)` if missing.
2. Build the model from the patched configs.
3. Load the projection + LoRA from `CKPT_DIR_PULLED`.
4. Iterate the **test** split, generate predictions, score every task.
5. Write `{LOCAL_RESULTS_DIR}/{RUN_ID}/predictions_*.json` + `metrics_summary.json`.
6. If `UPLOAD_RESULTS_TO_HF=True`, upload the folder to `{RUN_ID}/results/` on `HF_RUNS_REPO`.
"""))


CELLS.append(code("eval-run", """\
import shlex

RESULTS_DIR_LOCAL = WORK / LOCAL_RESULTS_DIR
RESULTS_DIR_LOCAL.mkdir(parents=True, exist_ok=True)

# evaluate.py forwards --checkpoint to utils.checkpoint.load_checkpoint,
# which reads <dir>/checkpoint_projection.pt + <dir>/checkpoint_lora/.
# Pass the DIRECTORY — passing the .pt file makes load_checkpoint mis-derive
# the filename (checkpoint_projection_projection.pt) and silently skip both
# projection AND LoRA, leaving you with raw 4-bit Vicuna.
CKPT_ARG = str(CKPT_DIR_PULLED)

extra = ''
if LLM_JUDGE:
    extra += ' --llm_judge'
    extra += f' --llm_judge_model {shlex.quote(LLM_JUDGE_MODEL)}'
    if LLM_JUDGE_BASE_URL:
        extra += f' --llm_judge_base_url {shlex.quote(LLM_JUDGE_BASE_URL)}'
    if LLM_JUDGE_MAX_SAMPLES:
        extra += f' --llm_judge_max_samples {int(LLM_JUDGE_MAX_SAMPLES)}'

if not UPLOAD_RESULTS_TO_HF:
    extra += ' --no_hf_upload'

# Metric comparability flags
extra += f' --bertscore_model {shlex.quote(BERTSCORE_MODEL)}'
extra += '' if BERTSCORE_RESCALE else ' --no-bertscore_rescale'
extra += f' --meteor_impl {METEOR_IMPL}'

print(f'Evaluating run_id    : {RUN_ID}')
print(f'Checkpoint           : {CKPT_ARG}')
print(f'Task                 : {TASK}   (split={SPLIT})')
print(f'Local results dir    : {RESULTS_DIR_LOCAL}/{RUN_ID}/')
print(f'Upload to HF         : {UPLOAD_RESULTS_TO_HF}')
print()

# NOTE: do NOT set HF_HUB_DISABLE_PROGRESS_BARS=1 here. On a cold cache the
# 4-bit Vicuna shard download is ~13GB and takes minutes on Colab T4 — hiding
# the bar makes the cell look frozen. TQDM_MININTERVAL=1.0 forces the per-task
# tqdm in run_inference to refresh every 1s.
!TRANSFORMERS_VERBOSITY=warning TOKENIZERS_PARALLELISM=false BITSANDBYTES_NOWELCOME=1 \\
PYTHONUNBUFFERED=1 TQDM_MININTERVAL=1.0 \\
python -u -m evaluation.evaluate \\
    --model_config    configs/model_config.yaml \\
    --train_config    configs/train_config.yaml \\
    --checkpoint      "{CKPT_ARG}" \\
    --run_id          "{RUN_ID}" \\
    --task            {TASK} \\
    --split           {SPLIT} \\
    --batch_size      {BATCH_SIZE} \\
    --max_new_tokens  {MAX_NEW_TOKENS} \\
    --output_dir      "{RESULTS_DIR_LOCAL}" \\
    --device          cuda{extra}
"""))


CELLS.append(md("eval-summary-md", """\
## 9. Show the metrics summary
"""))


CELLS.append(code("eval-summary", """\
import json as _json

_summary_path = RESULTS_DIR_LOCAL / RUN_ID / 'metrics_summary.json'
if not _summary_path.is_file():
    print(f'No metrics_summary.json at {_summary_path}. Did evaluate.py error out?')
else:
    summary = _json.loads(_summary_path.read_text())
    print(f'Dataset    : {summary.get("dataset_name")}')
    print(f'Run ID     : {summary.get("run_id")}')
    print(f'Split      : {summary.get("split")}')
    print()
    for task, metrics in (summary.get('metrics') or {}).items():
        print(f'─── {task.upper()} ───')
        for k, v in metrics.items():
            if isinstance(v, float):
                print(f'  {k:<22s}: {v:.4f}')
            else:
                print(f'  {k:<22s}: {v}')
        print()

    # List per-task prediction files for convenience
    print('Per-task prediction files:')
    for f in sorted((RESULTS_DIR_LOCAL / RUN_ID).glob('predictions_*.json')):
        n = len(_json.loads(f.read_text()))
        print(f'  {f.name}  ({n} samples)')
"""))


CELLS.append(md("eval-cleanup-md", """\
### Done

Final artifacts (also pushed to HF if `UPLOAD_RESULTS_TO_HF=True`):

```
{LOCAL_RESULTS_DIR}/{RUN_ID}/
  predictions_findings.json
  predictions_impression.json
  predictions_vqa.json       (MIMIC datasets only)
  metrics_summary.json
```

For free-form inference on individual images (with image preview, custom VQA questions, etc.) use the separate **`cxrvlm_colab_inference.ipynb`** notebook.
"""))


# ─────────────────────────────────────────────────────────────────────
# Write notebook
# ─────────────────────────────────────────────────────────────────────

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)")