#!/usr/bin/env python

import argparse
import json
import logging
import os
import sys
import time

import numpy as np
import SimpleITK as sitk
from tinydb import TinyDB, Query

from imaging.preprocessing import process_dataset
from config.constant import LABEL_MAP
from imaging.orientation import (best_symmetry_plane, best_upper_endplate_plane,
                                  segment_spinous_process, segment_vertebral_body,
                                  smooth_mask_sdf, measure_pedicle_width)
from visualization.res_bone_figure import (render_bone_figure,
                                           compute_normalizing_rotation,
                                           rotate_volume_to, rotated_grid,
                                           rotated_sitk_image_at,
                                           _rotate_plane_params,
                                           _shift_plane_params)
from imaging.transforms import merge_rotated_into_transform, level_file_path

data_root   = '/mnt/1220/Public/dataset/Spine/CTSpine1K/data/'
label_root  = '/mnt/1220/Public/dataset/Spine/CTSpine1K/label/'
output_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-2/'
output_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3/'

label_map = {
    'colon': 'conlon',
    'COVID-19': 'COVID-19',
    'HNSCC-3DCT-RT_neck': 'HNSCC-3DCT-RT_neck',
    'liver': 'Liver',
}

LUMBAR_LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5', 'L6')
MAX_Z_SPACING_MM = 4.0  # z spacing (mm) 大於此值者跳過整支 pipeline
MIN_LUMBAR_LEVELS = 2   # 符合 lumbar 的 label 少於此值者跳過整支 pipeline
MAX_PEDICLE_WIDTH_MM = 18.0  # 腰椎椎弓根寬度生理上限；量測超過即判為
                             # 橫突 / 肋在 SDF mask 內與弓根融合，記錄值
                             # clip 到此上限並標 _clipped（原始值仍記錄 _raw_mm）

LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs')
METADATA_DB = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                           'xfr_image_metadata.json')

logger = logging.getLogger('xfr_preprocess')


def _upsert_by_name(table, name, meta):
    """以 'name' 欄位為鍵 upsert（TinyDB 4.9 insert 不支援自訂 doc_id）。"""
    q = Query().name == name
    entry = dict(table.get(q) or {})
    entry.update(meta)
    entry['name'] = name
    if table.get(q) is None:
        table.insert(entry)
    else:
        table.update(entry, q)


def _looks_like_tinydb_table(v):
    """TinyDB table 資料格式：{數字字串 doc_id: doc}。"""
    return (isinstance(v, dict) and bool(v)
            and all(isinstance(k, str) and k.isdigit() for k in v))


def _migrate_legacy_metadata(path):
    """舊版純 JSON 檔（{影像 name: {...}}、無 TinyDB table 結構）：
    改名成 .legacy-<timestamp> 備份，內容匯入新 TinyDB。

    注意：TinyDB 4.x 落檔不含頂層 _default 鍵（只有 {table: {id: doc}}），
    舊判定「無 _default 即 legacy」會把正常 TinyDB 檔整體当成單一 legacy
    條目重寫（反覆包層、全檔損毀）。改以結構判定：頂層有任一
    {數字 id: doc} 形式的 table 即視為 TinyDB 檔、不做遷移。"""
    if not os.path.exists(path):
        return
    try:
        with open(path) as f:
            data = json.load(f)
    except (json.JSONDecodeError, OSError):
        return
    if not isinstance(data, dict) or not data:
        return
    if any(_looks_like_tinydb_table(v) for v in data.values()):
        return
    legacy = {k: v for k, v in data.items() if isinstance(v, dict)}
    if not legacy:
        return
    bak = f'{path}.legacy-{time.strftime("%Y%m%d_%H%M%S")}'
    os.replace(path, bak)
    print(f'Metadata db {path} is old pure-JSON format; renamed to {bak} '
          f'and migrating {len(legacy)} entrie(s)')
    t = TinyDB(path).table('images')
    for k, v in legacy.items():
        _upsert_by_name(t, k, dict(v))


class ImageMetadataDB:
    """影像 metadata 快取（TinyDB，table 'images'）：
    以影像檔名前綴（= 輸出資料夾名，如 1.3.6.1.4.1.9328.50.4.0001、
    liver_100、volume-covid19-A-0011_ct）為 key（存於文件 'name' 欄位），
    記錄 {'spacing': [x, y, z], 'labels': [label id]}。

    供 process_single_image 做跳過判定（z spacing / lumbar 層數），
    命中時不必讀影像 / label 檔。TinyDB 每次操作自動落檔，
    中途重跑可繼承。
    """

    def __init__(self, path=METADATA_DB):
        self.path = path
        _migrate_legacy_metadata(path)
        self._table = TinyDB(path).table('images')

    def get(self, name):
        return self._table.get(Query().name == name)

    def put(self, name, meta):
        _upsert_by_name(self._table, name, meta)

    def __len__(self):
        return len(self._table)


class _Tee:
    """同時輸出到 console 與 log 檔，逐行寫入。"""

    def __init__(self, console, log_fh):
        self.console = console
        self.log_fh = log_fh
        self.buf = ''

    def _emit(self, line):
        self.console.write(line + '\n')
        self.log_fh.write(line + '\n')

    def write(self, data):
        if not data:
            return
        self.buf += data
        while True:
            idx_n = self.buf.find('\n')
            idx_r = self.buf.find('\r')
            candidates = [i for i in (idx_n, idx_r) if i != -1]
            if not candidates:
                break
            idx = min(candidates)
            self._emit(self.buf[:idx])
            self.buf = self.buf[idx + 1:]

    def flush(self):
        if self.buf:
            self._emit(self.buf)
            self.buf = ''
        self.console.flush()
        self.log_fh.flush()

    def isatty(self):
        return False


def setup_tee(log_path):
    """把這個流程的 stdout/stderr 同時寫到 log_path（append、line-buffered）"""
    log_fh = open(log_path, 'a', buffering=1)
    sys.stdout = _Tee(sys.stdout, log_fh)
    sys.stderr = _Tee(sys.stderr, log_fh)


def _cortical_from_roi(roi_arr, bin_arr):
    """皮質遮罩（與舊未旋轉 _cortical.nii.gz 同定義）：
    門檻 = 骨頭 mask（bin_arr）內 CT 的 median HU；
    cortical = mask 內且 門檻 <= HU <= 10000。
    roi_arr / bin_arr 須同 grid（(z, y, x)）；回傳 uint8 0/1，
    輸入缺失或 mask 為空時回傳 None。"""
    if roi_arr is None or bin_arr is None:
        return None
    inside = bin_arr > 0
    if int(inside.sum()) == 0:
        return None
    threshold = float(np.median(roi_arr[inside]))
    return (inside & (roi_arr >= threshold) & (roi_arr <= 10000)).astype(np.uint8)


def _write_rotated_level(vol_dir, level, smd_path, mask_path, roi_path):
    """對單一 level：計算對齊旋轉（鏡稱面 -> +X、上終板 normal y=0），
    存到 <vol_dir>/rotated/（先旋轉 _binary_nn，再以它的非零 bbox（±4
    體素）作為所有旋轉輸出的裁切 grid，尺寸不求同未旋轉版、整顆骨頭
    保留，origin 隨裁切平移）：
      - _smd_resampled.nii.gz  三线性旋轉後 SMD（內負/外正）
      - _binary_sdf.nii.gz     旋轉 SMD 於 0.5 閾值（與未旋轉版一致）
      - _binary_nn.nii.gz      NN 旋轉後 0/1 遮罩（對比用，裁切基準）
      - _roi.nii.gz
      - _cortical.nii.gz       旋轉 CT 以骨頭 mask 內 median HU 為門檻
                               （取代舊的未旋轉 _cortical，定義相同）
    再畫「rotated」平面圖（bone + 平面 + VBODY [金] / 棘突 [紫] 著色），
    並在旋轉體上做 VBODY / 棘突分割存 label map
    （1=VBODY、2=棘突、3=other bone、0=background）。
    平面參數經 R 剛性旋轉並換元到輸出 grid 的局部座標。

    完成後把本 level 的旋轉座標（template / R / center / start / size）、
    SP / VBODY 分割兩平面與雙側椎弓根寬度（pedicle：L/R mm
    （mediolateral 中位數寬度，clip 上限 MAX_PEDICLE_WIDTH_MM）+ raw 值 +
    clipped 旗標、量測切片數、AP 窗）補寫進 <vol_dir>/transform.json
    （base 部分由 process_single_image 寫入）：
    供分割後把 rotated mask 反向映射回 data_root 原始 CT 座標
    （imaging/transforms.py）。

    旋轉來源 + 幾何模板：patient 未旋轉的 0.5mm _smd_resampled.nii.gz
    （新世代在 <vol>/crop/、舊世代在 <vol>/ 頂層；float、裁物件 bbox，
    SMD < 0.5 = 內部，與 _binary_sdf 同幾何）；
    缺失或格式不符時回退 _binary_nn.nii.gz（0.5mm 最近邻 0/1），
    此情形只產出 _binary_nn（+ roi / cortical / 圖 / label）。
    mask_path（未旋轉 0.5mm 0/1 遮罩，如 _binary_sdf）只供畫圖
    （render_bone_figure 以 uint8 讀取；float 檔會被讀成全 0），
    缺失時自動回退 _binary_nn。
    """
    volume_id = os.path.basename(vol_dir)

    template = None
    smd_arr = None
    if smd_path is not None and os.path.exists(smd_path):
        sa = sitk.GetArrayFromImage(sitk.ReadImage(smd_path)).astype(np.float32)
        if sa.size > 0 and np.isfinite(sa).all() and sa.min() < 0.0 and sa.max() > 0.0:
            template = sitk.ReadImage(smd_path)   # 0.5mm 幾何模板
            smd_arr = sa
        else:
            logger.warning(f'[rotated] {volume_id} {level}: _smd_resampled 不可用 '
                           f'(value)，改用 _binary_nn')
    if template is None:
        nn_in_path = level_file_path(vol_dir, level, 'binary_nn')
        if not os.path.exists(nn_in_path):
            logger.warning(f'[rotated] {volume_id} {level}: _smd_resampled / _binary_nn 皆不可用, skip')
            return
        template = sitk.ReadImage(nn_in_path)

    # 平面估計輸入：SMD 時取 0.5 閾值內部，與 _binary_sdf 同定義
    if smd_arr is not None:
        bin_arr = (smd_arr < 0.5).astype(np.uint8)
    else:
        bin_arr = (sitk.GetArrayFromImage(template) > 0).astype(np.uint8)
    if int(bin_arr.sum()) == 0:
        logger.warning(f'[rotated] {volume_id} {level}: empty binary, skip')
        return

    # 圖用遮罩須與模板同幾何（SDF 與 SMD 皆裁 bbox2；NN fallback 只有
    # NN 遮罩一致）；0.5mm 0/1，float 檔會被 uint8 讀取成全 0
    nn_in_path = level_file_path(vol_dir, level, 'binary_nn')
    if smd_arr is None:
        mask_path = nn_in_path if os.path.exists(nn_in_path) else mask_path
    if mask_path is None or not os.path.exists(mask_path):
        mask_path = nn_in_path if os.path.exists(nn_in_path) else None

    sym = best_symmetry_plane(bin_arr)
    symp = best_upper_endplate_plane(bin_arr)
    R, c_xyz = compute_normalizing_rotation(bin_arr, sym, symp)

    # 輸出 grid：未旋轉緊密 bbox 的 8 角點旋轉後的最小包圍盒 + margin——
    # 尺寸不受未旋轉版約束，保證旋轉後整顆骨頭保留（同尺寸旋轉切角落）
    start, size = rotated_grid(bin_arr.shape, R, c_xyz, margin=4)

    rotated_dir = os.path.join(vol_dir, 'rotated')
    os.makedirs(rotated_dir, exist_ok=True)

    # 先旋轉 _binary_nn（order=0，未旋轉 _binary_nn 先重採樣到模板 grid
    # 再旋轉），以其非零 bbox（外扩 4 體素）定所有旋轉輸出的裁切 grid
    # （fstart/fsize，模板 index 系）；NN 缺失 / 為空時回退緊密包圍盒
    rot_nn = None
    fstart = (int(start[0]), int(start[1]), int(start[2]))
    fsize = (int(size[0]), int(size[1]), int(size[2]))
    if os.path.exists(nn_in_path):
        nn_img = sitk.ReadImage(nn_in_path)
        nn_arr = (sitk.GetArrayFromImage(
            sitk.Resample(nn_img, template, interpolator=sitk.sitkNearestNeighbor,
                          defaultPixelValue=0)) > 0).astype(np.uint8)
        rot_nn_full = (rotate_volume_to(nn_arr, R, c_xyz, start, size,
                                        order=0, cval=0.0) > 0.5).astype(np.uint8)
        p_nn = os.path.join(rotated_dir, f'{level}_binary_nn.nii.gz')
        if int(rot_nn_full.sum()) > 0:
            zz, yy, xx = np.where(rot_nn_full > 0)
            m = 4
            z0 = max(0, int(zz.min()) - m)
            z1 = min(rot_nn_full.shape[0] - 1, int(zz.max()) + m)
            y0 = max(0, int(yy.min()) - m)
            y1 = min(rot_nn_full.shape[1] - 1, int(yy.max()) + m)
            x0 = max(0, int(xx.min()) - m)
            x1 = min(rot_nn_full.shape[2] - 1, int(xx.max()) + m)
            fstart = (fstart[0] + x0, fstart[1] + y0, fstart[2] + z0)
            fsize = (x1 - x0 + 1, y1 - y0 + 1, z1 - z0 + 1)
            rot_nn = rot_nn_full[z0:z1 + 1, y0:y1 + 1, x0:x1 + 1]
        else:
            rot_nn = rot_nn_full
        sitk.WriteImage(rotated_sitk_image_at(template, rot_nn, fstart), p_nn)
        logger.info(f'[rotated] saved {p_nn} (grid={list(fsize)}, start={fstart})')

    # 旋轉後 SMD：未旋轉 _smd_resampled 的 SMD 場經三线性旋轉得以保留
    # （float32 來源，map_coordinates 輸出維持浮點，<0.5 即為真閾值）；
    # 體外填充值取正 SMD（背景側），避免裁切邊界出現假的閾值穿越
    rot_sdf_bin = None
    if smd_arr is not None:
        rot_smd = rotate_volume_to(smd_arr, R, c_xyz, fstart, fsize,
                                   order=1, cval=float(smd_arr.max()))
        p_smd = os.path.join(rotated_dir, f'{level}_smd_resampled.nii.gz')
        sitk.WriteImage(rotated_sitk_image_at(template, rot_smd.astype(np.float32),
                                              fstart), p_smd)
        logger.info(f'[rotated] saved {p_smd}')

        # _binary_sdf.nii.gz：旋轉 SMD 於 0.5 閾值（與未旋轉版一致）
        rot_sdf_bin = (rot_smd < 0.5).astype(np.uint8)
        p_sdf = os.path.join(rotated_dir, f'{level}_binary_sdf.nii.gz')
        sitk.WriteImage(rotated_sitk_image_at(template, rot_sdf_bin, fstart), p_sdf)
        logger.info(f'[rotated] saved {p_sdf}')

    # 旋轉後的 roi（三线性）：先重採樣到模板 grid，再旋轉到裁切 grid
    roi_arr = None
    rot_roi_arr = None
    if roi_path is not None and os.path.exists(roi_path):
        roi_img = sitk.ReadImage(roi_path)
        roi_arr = sitk.GetArrayFromImage(
            sitk.Resample(roi_img, template, defaultPixelValue=0.0))
        rot_roi = rotate_volume_to(roi_arr, R, c_xyz, fstart, fsize,
                                   order=1, cval=0.0)
        p_roi = os.path.join(rotated_dir, f'{level}_roi.nii.gz')
        sitk.WriteImage(rotated_sitk_image_at(template, rot_roi, fstart), p_roi)
        logger.info(f'[rotated] saved {p_roi}')
        rot_roi_arr = rot_roi

    # 分割 / cortical 輸入：SMD 時直接用旋轉 _binary_sdf（0.5 閾值、
    # 次體素平滑）；NN fallback 時以 smooth_mask_sdf 記憶體內平滑（不存檔）
    if rot_sdf_bin is not None:
        rot_bin_seg = rot_sdf_bin
    else:
        rot_bin_seg = smooth_mask_sdf(rot_nn, sigma=1.5).astype(np.uint8)

    # 旋轉後的 cortical：旋轉 CT 以骨頭 mask 內 median HU 為門檻
    #（取代舊的未旋轉 _cortical.nii.gz，定義相同）
    rot_cort = _cortical_from_roi(rot_roi_arr, rot_bin_seg)
    if rot_cort is not None:
        p_cort = os.path.join(rotated_dir, f'{level}_cortical.nii.gz')
        sitk.WriteImage(rotated_sitk_image_at(template, rot_cort, fstart), p_cort)
        logger.info(f'[rotated] saved {p_cort}')
    else:
        logger.warning(f'[rotated] {volume_id} {level}: 無旋轉 CT / mask，跳過 _cortical')

    # 用旋轉後的平面畫圖（rotated 版 planes），並畫出 VBODY / 棘突著色
    #（VBODY=金、棘突=紫，與 label map 對應）；皮質著色由未旋轉 CT + mask
    # 現算（未旋轉 _cortical 不再存檔）。shift=fstart 讓點雲 / 平面落在
    # 存檔 rotated/*.nii.gz 的局部 index 系——與 CBT 螺絲圖同座標系
    p_fig = os.path.join(rotated_dir, f'{level}_planes.png')
    if mask_path is not None:
        fig_cortical = _cortical_from_roi(roi_arr, bin_arr)
        fig = render_bone_figure(volume_id, level, mask_path, fig_cortical,
                                 planes_only=False,
                                 rotation=(R, c_xyz, fstart), output_path=p_fig)
        if fig is not None:
            logger.info(f'[rotated] saved {fig}')
    else:
        logger.warning(f'[rotated] {volume_id} {level}: 無可用 0/1 遮罩，跳過平面圖')

    # 旋轉體上的 VBODY / 棘突分割：平面參數隨 R 剛性旋轉（rotated 體與
    # 原始體只差一個剛性變換，平面隨動即可，不需重新偵測），
    # label map 存同資料夾：1 = VBODY（椎體），2 = 棘突（spinous process），
    # 3 = other bone（其餘骨頭），0 = background。
    # 換元到輸出 grid 的局部座標（origin 在模板 index fstart 處）
    sym_rot = _shift_plane_params(_rotate_plane_params(sym, R, c_xyz), fstart)
    symp_rot = (_shift_plane_params(_rotate_plane_params(symp, R, c_xyz), fstart)
                if symp is not None else None)
    sp_mask, sp_th, sp_info = segment_spinous_process(rot_bin_seg, sym_rot)
    vb_mask, vb_th, vb_info = segment_vertebral_body(rot_bin_seg, sym_rot, symp_rot,
                                                     sp_th, sp_info['mode'])
    label_arr = np.zeros(rot_bin_seg.shape, dtype=np.uint8)
    label_arr[rot_bin_seg > 0] = 3                       # other bone
    if vb_mask is not None:
        label_arr[vb_mask] = 1                        # VBODY
    if sp_mask is not None:
        label_arr[sp_mask] = 2                        # spinous process
    p_lbl = os.path.join(rotated_dir, f'{level}_label.nii.gz')
    sitk.WriteImage(rotated_sitk_image_at(template, label_arr, fstart), p_lbl)
    logger.info(f'[rotated] saved {p_lbl} '
                f'(vbody={int((label_arr == 1).sum())}, spinous={int((label_arr == 2).sum())}, '
                f'other={int((label_arr == 3).sum())}, '
                f'sp_mode={sp_info["mode"]}, vb_mode={vb_info["mode"]})')

    # 雙側椎弓根寬度（左右 mediolateral，rotated 體上量測）：椎管 AP 窗內、
    # 中線帶外椎弓根區的每側最靠中線 2D 分塊、逐切片 x 向跨度沿 z 取中位數
    #（與上方 SP / VBODY 分割同用 rot_bin_seg / sym_rot，ap / s 定義一致）
    ped = measure_pedicle_width(rot_bin_seg, sym_rot, sp_info, vb_info)
    sp_x = float(template.GetSpacing()[0])

    def _ped_mm(res):
        """(記錄值 mm（clip 上限）, 原始值 mm, 是否被 clip)。"""
        raw = res['width_vox'] * sp_x if res['width_vox'] is not None else None
        clipped = raw is not None and raw > MAX_PEDICLE_WIDTH_MM
        val = (min(raw, MAX_PEDICLE_WIDTH_MM) if raw is not None else None)
        return (round(val, 2) if val is not None else None,
                round(raw, 2) if raw is not None else None, bool(clipped))
    L_rec, L_raw, L_clip = _ped_mm(ped['L'])
    R_rec, R_raw, R_clip = _ped_mm(ped['R'])
    ped_entry = {
        'L_mm': L_rec, 'R_mm': R_rec,
        'L_raw_mm': L_raw, 'R_raw_mm': R_raw,
        'L_clipped': L_clip, 'R_clipped': R_clip,
        'L_slices': ped['L']['n_slices'], 'R_slices': ped['R']['n_slices'],
        'L_z': ped['L']['z'], 'R_z': ped['R']['z'],
        'ap_window': (None if ped['ap_front'] is None
                      else [round(ped['ap_front'], 1), round(ped['ap_back'], 1)]),
        'band_w': (None if ped['band_w'] is None else round(ped['band_w'], 2)),
        'reason': ped['reason'],
    }

    def _fmt_mm(v):
        return 'n/a' if v is None else f'{v:.2f}'
    clip_note = (f' (raw>{MAX_PEDICLE_WIDTH_MM:g} 判橫突 / 肋融合, clipped)'
                 if (ped_entry['L_clipped'] or ped_entry['R_clipped']) else '')
    logger.info('[rotated] %s %s: pedicle width L=%s, R=%s mm '
                '(raw L=%s, R=%s%s; slices L=%d, R=%d; %s)',
                volume_id, level, _fmt_mm(ped_entry['L_mm']),
                _fmt_mm(ped_entry['R_mm']), _fmt_mm(ped_entry['L_raw_mm']),
                _fmt_mm(ped_entry['R_raw_mm']), clip_note,
                ped_entry['L_slices'], ped_entry['R_slices'], ped_entry['reason'])

    # 記錄 rotated 座標鏈（template 選擇 / R / center / start）+ 本 level
    # SP / VBODY 分割所用的兩平面（rotated/ 局部系）+ 雙側椎弓根寬度
    #（pedicle，量測於同 rot_bin_seg），補進 transform.json；
    # rotated disk a -> 未旋轉 template disk:
    # W = R^T (a + start - center) + center（template disk 系，standardize 後）。
    # 平面記錄後 core/optimizer 的 xfr 管線直接讀用、不在影像上重新估算
    # （imaging.transforms.load_rotated_planes）；sym_rot / symp_rot 即本函式
    # 上方 _label map / 平面圖同用的換系平面。
    try:
        merge_rotated_into_transform(vol_dir, level, {
            'template': 'smd_resampled' if smd_arr is not None else 'binary_nn',
            'R': [[float(v) for v in row] for row in R],
            'center': [float(v) for v in c_xyz],
            'start': [int(v) for v in fstart],
            'size': [int(v) for v in fsize],
            'sym': sym_rot,
            'endplate': symp_rot,
            'pedicle': ped_entry,
        })
        logger.info(f'[rotated] transform.json: recorded {level} rotated section '
                    f'(template={ "smd_resampled" if smd_arr is not None else "binary_nn" }, '
                    f'start={list(fstart)})')
    except Exception as e:
        logger.warning(f'[rotated] {volume_id} {level}: transform.json 記錄失敗（{e}）；'
                       f'該 level 需 --rebuild 才能反向映射')


def make_lumbar_post_process():
    """每個 volume 處理完後，對其 lumbar level：
    1) 畫「骨頭 + 方向平面」圖（不畫螺絲、不做棘突 / 椎體分割）
       -> <volume_dir>/crop/<level>_planes.png（未旋轉檔同在 crop/）
    2) 計算對齊旋轉，存旋轉後的 smd_resampled / binary_sdf / binary_nn / roi
       + cortical + 旋轉平面圖（含 VBODY / 棘突著色）+ label map -> <volume_dir>/rotated/

    post_process 由 process_dataset 呼叫：(volume_dir, processed_labels)。
    processed_labels 為該 volume 實際存在的 label id（int），對照 LABEL_MAP。
    """

    def _post_process(vol_dir, processed_labels):
        volume_id = os.path.basename(vol_dir)
        crop_dir = os.path.join(vol_dir, 'crop')
        for n in processed_labels:
            level = LABEL_MAP.get(int(n))
            if level not in LUMBAR_LEVELS:
                continue
            # 未旋轉輸入：新世代在 crop/、舊世代在 <vol>/ 頂層
            smd_res_path = level_file_path(vol_dir, level, 'smd_resampled')
            nn_path = level_file_path(vol_dir, level, 'binary_nn')
            if not os.path.exists(smd_res_path) and not os.path.exists(nn_path):
                continue
            # 畫圖用 0/1 遮罩：優 _binary_sdf（0.5mm 平滑），缺則 _binary_nn
            sdf_path = level_file_path(vol_dir, level, 'binary_sdf')
            mask_path = sdf_path if os.path.exists(sdf_path) else nn_path
            roi_path = level_file_path(vol_dir, level, 'roi')

            # 1) 原始（未旋轉）planes 圖；皮質著色由未旋轉 CT + mask 現算
            #（未旋轉 _cortical 不再存檔）；圖與未旋轉檔同在 crop/
            os.makedirs(crop_dir, exist_ok=True)
            output_path = os.path.join(crop_dir, f'{level}_planes.png')
            if os.path.exists(mask_path):
                fig_cortical = None
                if os.path.exists(roi_path):
                    mask_img = sitk.ReadImage(mask_path)
                    roi_img = sitk.ReadImage(roi_path)
                    roi_arr = sitk.GetArrayFromImage(
                        sitk.Resample(roi_img, mask_img, defaultPixelValue=0.0))
                    fig_cortical = _cortical_from_roi(
                        roi_arr, sitk.GetArrayFromImage(mask_img))
                path = render_bone_figure(volume_id, level, mask_path, fig_cortical,
                                          planes_only=True, output_path=output_path)
                if path is not None:
                    logger.info(f'[crop] saved {path}')

            # 2) 旋轉對齊：rotated/ 的 smd_resampled + binary_sdf + binary_nn
            #    + roi + cortical + planes 圖（含 VBODY / 棘突著色）+ label
            _write_rotated_level(vol_dir, level, smd_res_path, mask_path, roi_path)

    return _post_process


def main():
    parser = argparse.ArgumentParser(description='Preprocess CT spine dataset.')
    parser.add_argument('--max-images', type=int, default=None, dest='max_images',
                        help='Process at most this number of images per dataset (default: all).')
    parser.add_argument('--output-dir', type=str, default=None, dest='output_dir',
                        help='Override the default output dir (e.g. repair run on an '
                             'older generation). Default: the module-level output_dir.')
    args = parser.parse_args()
    # local alias（避免 rebind module-level output_dir 造成 UnboundLocalError）
    out_dir = args.output_dir if args.output_dir is not None else output_dir

    # log 檔（console 與檔案同時輸出）
    os.makedirs(LOG_DIR, exist_ok=True)
    log_path = os.path.join(LOG_DIR, f'xfr_preprocess_{time.strftime("%Y%m%d_%H%M%S")}.log')
    setup_tee(log_path)
    # basicConfig 在 setup_tee 之後，讓 handler 寫入 Tee（同時進 console 與 log 檔）
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    logger.info(f'Log file: {log_path}')
    logger.info(f'Command: {sys.executable} {" ".join(sys.argv)}')
    logger.info(f'Working directory: {os.getcwd()}')
    logger.info(f'Output dir: {out_dir}')

    # metadata db：跳過判定（z spacing / lumbar 層數）命中時免讀影像 / label 檔
    metadata_db = ImageMetadataDB()
    logger.info(f'Metadata db: {metadata_db.path} ({len(metadata_db)} entries)')

    post_process = make_lumbar_post_process()

    for key, value in label_map.items():
        data_dir = os.path.join(data_root, key)
        label_dir = os.path.join(label_root, value)
        process_dataset(data_dir, label_dir, out_dir, max_images=args.max_images,
                        post_process=post_process, max_z_spacing=MAX_Z_SPACING_MM,
                        allowed_levels=LUMBAR_LEVELS, min_levels=MIN_LUMBAR_LEVELS,
                        metadata_cache=metadata_db, crop_subdir=True)

if __name__ == '__main__':
    main()