#!/home/xfr/.conda/envs/cbt/bin/python

"""把每個 lumbar level 的 rotated/{L}_label.nii.gz（旋轉座標，
1=VBODY 椎體、2=SP 棘突、3=other bone）反旋轉回原始座標，
把所有 level 合併成「每 patient 一個」.nii.gz，存到 CTSpine1K
原始 dataset 下面（座標系與 data/ 裡的原始影像檔一致）。

做法：
  1) 依 level 重算對齊旋轉 (R, c_xyz)（與 _write_rotated_level 完全
     同路徑：_smd_resampled 有效時以 <0.5 取骨頭 mask，否則 _binary_nn
     以 >0），再從旋轉 label 檔與未旋轉模板檔的 origin 差回復裁切
     起點 fstart（(O_rot - O_tpl) 換元、必須為整數）。
  2) NN 反向映射把 label 映回未旋轉 0.5mm 模板 grid。旋轉是 index 空間
     剛性變換（物理座標系統只平移 fstart），所以換元後只是單純的
     R 旋轉，不需要重新估算任何平面。
  3) std（0.5mm 標準化空間）-> 原始 CT 物理座標的對角仿射 mapping：
     standardize_affine 每翻一個軸就把整份資料平移到世界原點的鏡射面
     （b_i = -(2O_i + d_i·0.5·(L_i-1))，翻軸條件 x: d_x=+1、y: d_y=+1、
     z: d_z=-1），ap_flip（metadata db）再令 y 軸鏡射。全部參數都可從
     原始 CT 幾何 + ap_flip 算出，見 _std_to_orig_affine。
4) NN 重取樣到原始 CT grid（/mnt/1220/Public/dataset/Spine/CTSpine1K/
      data/<子目錄>/<name>.nii.gz 的幾何），依 level 重編號後合併：
      Lx 的 VBODY = 2x-1、SP = 2x（L1..L6 -> 1..12），0 = 背景；
      rotated label 的 3（other bone）不納入。
   每 level 驗證重算的 (R, c)：主要用正向 identity —— 以 _write_rotated_level
      完全同參數重旋轉未旋轉 check 檔（_smd_resampled 三線性、缺則
      _binary_nn NN），與存檔旋轉檔逐體素比對，mismatch > 1e-4 跳過
      （精確驗證 (R,c) 與該輸出生成版本一致，不受薄結構影響）；無可比檔
      時 fallback 用 round-trip 包含率（旋轉 check 檔反向旋轉後與未旋轉
      bone mask 互相包含於 2-voxel 膨脹內，min 方向 < --min-contain
      跳過該 level）。
   寫檔前對位檢查：合併 label 落在原始 CT 體內（HU > -100）比例 < 0.98
   或 median HU < 100 時判定錯位、跳過。

輸出：<dest>/<子目錄>/<name>.nii.gz（uint8），預設 dest =
/mnt/1220/Public/dataset/Spine/CTSpine1K/label_vbody_sp/，
檔名與原始影像檔相同。已存在的輸出預設跳過（--force 覆蓋）。
"""

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

import numpy as np
import SimpleITK as sitk
from scipy.ndimage import map_coordinates

from imaging.orientation import (best_symmetry_plane, best_upper_endplate_plane)
from imaging.transforms import level_file_path
from visualization.res_bone_figure import (compute_normalizing_rotation,
                                           rotated_grid, rotate_volume_to)

DATA_ROOT = '/mnt/1220/Public/dataset/Spine/CTSpine1K/data/'
DATASET_ROOT = '/mnt/1220/Public/dataset/Spine/CTSpine1K/'
DEFAULT_DEST = os.path.join(DATASET_ROOT, 'label_vbody_sp')
# 預設依序掃的世代資料夾（先新後舊；同名 volume 以先處理者為準）
DEFAULT_ROOTS = [
    '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3',
    '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-2',
]
LUMBAR_LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5', 'L6')

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_orig_labels')


def load_ap_flip():
    """metadata db（TinyDB 純 JSON：{'images': {doc_id: doc}}）的
    name -> ap_flip（bool）；db 缺失 / 解析失敗時回 {}。"""
    try:
        with open(METADATA_DB) as f:
            data = json.load(f)
    except (json.JSONDecodeError, OSError) as e:
        logger.warning(f'metadata db {METADATA_DB} 讀取失敗：{e}')
        return {}
    tab = data.get('images') if isinstance(data, dict) else None
    if not isinstance(tab, dict):
        return {}
    return {v['name']: bool(v.get('ap_flip', False))
            for v in tab.values() if isinstance(v, dict) and 'name' in v}


# ---------------------------------------------------------------- log 設定
def setup_tee(log_path):
    """console 與 log 檔同時輸出（append、line-buffered）。"""
    class _Tee:
        def __init__(self, console, fh):
            self.console, self.fh, self.buf = console, fh, ''

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

        def flush(self):
            if self.buf:
                line, self.buf = self.buf, ''
                self.console.write(line + '\n')
                self.fh.write(line + '\n')
            self.console.flush()
            self.fh.flush()

    fh = open(log_path, 'a', buffering=1)
    sys.stdout = _Tee(sys.stdout, fh)
    sys.stderr = _Tee(sys.stderr, fh)


# ---------------------------------------------------------------- 幾何工具
def _axis_vectors(img):
    """direction 為純對角（±1）時回傳 (signs (x,y,z), spacing)，否則 None。"""
    d = np.asarray(img.GetDirection(), dtype=float).reshape(3, 3)
    diag = np.diag(d)
    if not (np.allclose(d, np.diag(diag), atol=1e-6)
            and np.allclose(np.abs(diag), 1.0, atol=1e-6)):
        return None
    return diag, np.asarray(img.GetSpacing(), dtype=float)


def _fstart_from_geometries(rot_img, t_img):
    """從旋轉檔與未旋轉模板檔的 origin 差回復 fstart（模板 index 系，
    (x,y,z) 整數）；幾何不一致（非對角 direction、spacing / direction
    不同）或差不是整數時回傳 None。"""
    ax_r, ax_t = _axis_vectors(rot_img), _axis_vectors(t_img)
    if ax_r is None or ax_t is None:
        return None
    (d_r, s_r), (d_t, s_t) = ax_r, ax_t
    if not (np.allclose(d_r, d_t, atol=1e-6) and np.allclose(s_r, s_t, atol=1e-6)):
        return None
    o_r = np.asarray(rot_img.GetOrigin(), dtype=float)
    o_t = np.asarray(t_img.GetOrigin(), dtype=float)
    fs = (o_r - o_t) / (d_t * s_t)
    if not np.allclose(fs, np.round(fs), atol=1e-6):
        return None
    return np.round(fs).astype(int)


def _std_to_orig_affine(orig_img, t_img, ap_flip):
    """std（0.5mm 裁切檔、standardize_affine 後、ap_flip 後）-> 原始 CT
    物理座標（LPS）：P_orig = A·P_std + b。回傳 (A (3,) 對角, b (3,))；
    原始 direction 非純對角時回傳 None。t_img = 該 level 的未旋轉標準化
    裁切檔（disk 上的後 standardize 幾何）。

    standardize_affine 是【逐檔（每 level 裁切）】處理：翻軸（nibabel：
    affine 對角與平移同反號 + 資料 flip）在 worlds 中是整份資料平移
    t_i = 2·O_post_i - d_i·h·(L_i-1)（O_post/L 為 disk 上該檔的
    origin/尺寸、d 為原始 direction、h=0.5；翻軸條件 RAS 對角 < 0
    = x: d_x=+1、y: d_y=+1、z: d_z=-1），故 b_i = -t_i（翻軸的軸）。
    ap_flip 是 y 軸純資料鏡射（幾何不變，作用在整張 0.5mm full grid）：
    y_ap = C - y_orig，C = 2·O_full_y + d_y·h·(L_full_y-1)（full grid
    幾何 = 原始 CT），發生在 standardize 之前。"""
    d_full = np.asarray(orig_img.GetDirection(), dtype=float).reshape(3, 3)
    d = np.diag(d_full)
    if not (np.allclose(d_full, np.diag(d), atol=1e-6)
            and np.allclose(np.abs(d), 1.0, atol=1e-6)):
        return None
    h = 0.5
    O_full = np.asarray(orig_img.GetOrigin(), dtype=float)
    s = np.asarray(orig_img.GetSpacing(), dtype=float)
    n = np.asarray(orig_img.GetSize(), dtype=float)           # (x, y, z)
    L_full = np.ceil(n * s / h - 1e-6)                        # 0.5mm full grid
    O_post = np.asarray(t_img.GetOrigin(), dtype=float)
    L = np.asarray(t_img.GetSize(), dtype=float)              # 裁切檔尺寸
    flips = (d[0] == 1.0, d[1] == 1.0, d[2] == -1.0)
    A = np.ones(3)
    b = np.zeros(3)
    for i in range(3):
        if flips[i]:
            b[i] = -(2.0 * O_post[i] - d[i] * h * (L[i] - 1.0))
    if ap_flip:
        # y0 --AP 鏡射--> y1 = C - y0 --std 平移--> y_std = y1 + t_y
        # => y0 = -y_std + (C + t_y)
        A[1] = -1.0
        C_full = 2.0 * O_full[1] + d[1] * h * (L_full[1] - 1.0)
        if flips[1]:
            t_y = 2.0 * O_post[1] - d[1] * h * (L[1] - 1.0)
            b[1] = C_full + t_y
        else:
            b[1] = C_full
    return A, b


def _affine_sitk(A, b):
    """對角 A + 平移 b 的 sitk 仿射（LPS 物理 -> LPS 物理）。
    p' = A·p + b；SetMatrix 為 row-major 3x3、SetTranslation 獨立。
    自檢三軸 + 原點。"""
    tr = sitk.AffineTransform(3)
    tr.SetMatrix([float(A[0]), 0.0, 0.0,
                  0.0, float(A[1]), 0.0,
                  0.0, 0.0, float(A[2])])
    tr.SetTranslation([float(b[0]), float(b[1]), float(b[2])])
    for i in range(3):
        p = [0.0, 0.0, 0.0]
        p[i] = 123.0
        got = np.array(tr.TransformPoint(tuple(p)))
        exp = np.array(b, dtype=float)
        exp[i] = A[i] * 123.0 + b[i]
        if not np.allclose(got, exp, atol=1e-6):
            raise RuntimeError(f'affine self-check failed: axis {i} '
                               f'{got} != {exp}')
    return tr


def _rotate_back(arr, R, c_xyz, fstart, shape_zyx, order=0, cval=0.0):
    """rotated grid 上的 field 反向旋轉到未旋轉模板 grid（shape_zyx）。
    rotated 檔 array 座標 a = 模板座標 V - fstart；反向映射
    V = R(W - c) + c，W 為模板 grid 上每格（x,y,z 向量，同
    compute_normalizing_rotation 慣例）。order / cval 同 map_coordinates。"""
    nz, ny, nx = shape_zyx
    idx = np.indices((nz, ny, nx), dtype=np.float64)
    W = np.stack([idx[2], idx[1], idx[0]], axis=0)          # (3, z, y, x)
    c = np.asarray(c_xyz, dtype=np.float64)[:, None, None, None]
    Rf = np.asarray(R, dtype=np.float64)
    V = np.tensordot(Rf, W - c, axes=([1], [0])) + c
    fs = np.asarray(fstart, dtype=np.float64)[:, None, None, None]
    a = V - fs                                              # rotated 檔 array 座標
    return map_coordinates(arr, [a[2], a[1], a[0]],
                           order=order, cval=cval, mode='constant')


# ---------------------------------------------------------------- level 處理
def _load_unrotated_template(vol_dir, level):
    """_write_rotated_level 的旋轉來源：_smd_resampled（有效時）否則
    _binary_nn。回傳 (img, arr, is_smd)；皆無則 None。
    未旋轉檔位置經 level_file_path（新世代 crop/、舊世代頂層）。"""
    smd_path = level_file_path(vol_dir, level, 'smd_resampled')
    if 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:
            return sitk.ReadImage(smd_path), sa, True
    nn_path = level_file_path(vol_dir, level, 'binary_nn')
    if os.path.exists(nn_path):
        na = sitk.GetArrayFromImage(sitk.ReadImage(nn_path))
        if na.size > 0 and int((na > 0).sum()) > 0:
            return sitk.ReadImage(nn_path), na, False
    return None


def _bone_mask(t_arr, is_smd):
    return (t_arr < 0.5) if is_smd else (t_arr > 0)


def _forward_identity_check(vol_dir, level, R, c_xyz, t_img, t_arr, is_smd):
    """以 _write_rotated_level 完全同參數（fstart/fsize 由未旋轉 _binary_nn
    旋轉後 bbox±4、order / cval 同 writer）重旋轉、與存檔的旋轉檔逐體素
    比對，回傳 mismatch 比例（(R,c) 相同時 ≈0，精確驗證、不受物件厚度
    影響）；無可比檔（缺檔 / 形狀不符）時回 None。"""
    rotated_dir = os.path.join(vol_dir, 'rotated')
    if is_smd:
        target = os.path.join(rotated_dir, f'{level}_smd_resampled.nii.gz')
        src_arr, order, cval = t_arr, 1, float(t_arr.max())
    else:
        target = os.path.join(rotated_dir, f'{level}_binary_nn.nii.gz')
        src_arr, order, cval = (t_arr > 0).astype(np.uint8), 0, 0.0
    if not os.path.exists(target):
        return None
    start, size = rotated_grid(t_arr.shape, R, c_xyz, margin=4)
    fstart = (int(start[0]), int(start[1]), int(start[2]))
    fsize = (int(size[0]), int(size[1]), int(size[2]))
    nn_path = level_file_path(vol_dir, level, 'binary_nn')
    if os.path.exists(nn_path):
        nn_img = sitk.ReadImage(nn_path)
        nn_arr = (sitk.GetArrayFromImage(
            sitk.Resample(nn_img, t_img, interpolator=sitk.sitkNearestNeighbor,
                          defaultPixelValue=0)) > 0).astype(np.uint8)
        rot_nn_full = (rotate_volume_to(nn_arr, R, c_xyz, fstart, fsize,
                                        order=0, cval=0.0) > 0.5).astype(np.uint8)
        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)
    out = rotate_volume_to(src_arr, R, c_xyz, fstart, fsize, order=order,
                           cval=cval)
    saved = sitk.GetArrayFromImage(sitk.ReadImage(target))
    if out.shape != saved.shape:
        return None
    if order == 0:
        return float(((out > 0.5) != (saved > 0.5)).mean())
    return float((np.abs(out.astype(np.float64)
                         - saved.astype(np.float64)) > 0.02).mean())


def _containment_roundtrip(vol_dir, level, R, c_xyz, fstart, t_arr, is_smd,
                           it=2):
    """fallback 驗證：旋轉 check 檔反向旋轉後與未旋轉 bone mask 互相包含於
    it-voxel 膨脹內（薄結構的 round-trip IoU 天生低，用包含率代替）。
    回傳 min 方向包含率；無 check 檔時回 None。"""
    from scipy.ndimage import binary_dilation
    rotated_dir = os.path.join(vol_dir, 'rotated')
    candidates = [
        (f'{level}_smd_resampled.nii.gz', 1),
        (f'{level}_binary_sdf.nii.gz', 0),
        (f'{level}_binary_nn.nii.gz', 0),
    ]
    if not is_smd:
        candidates = candidates[1:]
    B0 = _bone_mask(t_arr, is_smd)
    for fname, order in candidates:
        p = os.path.join(rotated_dir, fname)
        if not os.path.exists(p):
            continue
        a = sitk.GetArrayFromImage(sitk.ReadImage(p)).astype(np.float64)
        if a.size == 0:
            continue
        cval = float(a.max()) if fname.endswith('_smd_resampled.nii.gz') else 0.0
        b = _rotate_back(a, R, c_xyz, fstart, t_arr.shape, order=order,
                         cval=cval)
        B1 = b < 0.5 if fname.endswith('_smd_resampled.nii.gz') else b > 0.5
        if not B1.any() or not B0.any():
            continue
        da = binary_dilation(B1, iterations=it)
        db = binary_dilation(B0, iterations=it)
        return float(min((B1 & db).sum() / B1.sum(), (B0 & da).sum() / B0.sum()))
    return None


def process_level(vol_dir, level, min_contain):
    """單一 level：重算 (R, c)、回復 fstart、驗證、反向旋轉 label。
    回傳 (unrotated_label_arr (z,y,x) uint8 0/1/2/3, t_img, info)；失敗回 None。"""
    name = os.path.basename(vol_dir)
    rot_lbl_path = os.path.join(vol_dir, 'rotated', f'{level}_label.nii.gz')
    if not os.path.exists(rot_lbl_path):
        return None
    rot_img = sitk.ReadImage(rot_lbl_path)
    rot_arr = sitk.GetArrayFromImage(rot_img).astype(np.uint8)

    tpl = _load_unrotated_template(vol_dir, level)
    if tpl is None:
        logger.warning(f'{name} {level}: 無可用未旋轉模板 '
                       f'（_smd_resampled / _binary_nn），跳過')
        return None
    t_img, t_arr, is_smd = tpl

    fstart = _fstart_from_geometries(rot_img, t_img)
    if fstart is None:
        logger.warning(f'{name} {level}: 旋轉 / 未旋轉幾何不一致，跳過')
        return None

    bin_arr = _bone_mask(t_arr, is_smd).astype(np.uint8)
    sym = best_symmetry_plane(bin_arr)
    symp = best_upper_endplate_plane(bin_arr)
    R, c_xyz = compute_normalizing_rotation(bin_arr, sym, symp)

    # 驗證重算的 (R, c)：主要用 forward identity（與生成時同參數重旋轉、
    # 逐體素比對存檔；(R,c) 相同時 ≈0）；無可比檔時 fallback 用
    # round-trip 包含率（min_contain 門檻）。
    fwd = _forward_identity_check(vol_dir, level, R, c_xyz, t_img, t_arr, is_smd)
    contain = None
    if fwd is not None:
        if fwd > 1e-4:
            logger.warning(f'{name} {level}: 正向重旋轉 mismatch {fwd:.4f} '
                           f'> 1e-4，(R,c) 與生成參數不符，跳過')
            return None
    else:
        contain = _containment_roundtrip(vol_dir, level, R, c_xyz, fstart,
                                         t_arr, is_smd)
        if contain is None:
            logger.warning(f'{name} {level}: 無 check 檔可驗證，跳過')
            return None
        if contain < min_contain:
            logger.warning(f'{name} {level}: round-trip 包含率 {contain:.4f} '
                           f'< {min_contain}，跳過該 level')
            return None

    unrot = _rotate_back(rot_arr.astype(np.float64), R, c_xyz, fstart,
                         t_arr.shape).astype(np.uint8)
    n_lbl = int((unrot > 0).sum())
    B0 = _bone_mask(t_arr, is_smd)
    bone_frac = (int(((unrot > 0) & B0).sum()) / n_lbl) if n_lbl else 0.0
    info = {
        'forward_mismatch': fwd,
        'containment': contain,
        'fstart': [int(v) for v in fstart],
        'bone_frac': bone_frac,
        'rotated': {'vbody': int((rot_arr == 1).sum()),
                    'sp': int((rot_arr == 2).sum()),
                    'other': int((rot_arr == 3).sum())},
        'unrotated': {'vbody': int((unrot == 1).sum()),
                      'sp': int((unrot == 2).sum()),
                      'other': int((unrot == 3).sum())},
    }
    return unrot, t_img, info


def _resample_nearest(ref_img, src_img, transform=None):
    r = sitk.ResampleImageFilter()
    r.SetReferenceImage(ref_img)
    r.SetInterpolator(sitk.sitkNearestNeighbor)
    r.SetDefaultPixelValue(0)
    if transform is not None:
        r.SetTransform(transform)
    return r.Execute(src_img)


def process_volume(vol_dir, orig_img, key, dest_dir, min_contain, force,
                   ap_flip):
    """所有 level 反旋轉 + 合併 + 重取樣到原始 CT grid，存「一個」.nii.gz。
    回傳 (out_path, status, {level: info}, vol_info)；失敗時 out_path=None、
    status 為原因。"""
    name = os.path.basename(vol_dir)
    nxi, nyi, nzi = orig_img.GetSize()          # sitk GetSize = (x, y, z)
    merged = np.zeros((nzi, nyi, nxi), dtype=np.uint8)   # array = (z, y, x)
    level_info = {}
    maps = []

    for li, level in enumerate(LUMBAR_LEVELS):
        res = process_level(vol_dir, level, min_contain)
        if res is None:
            continue
        unrot, t_img, info = res
        # std->orig mapping 依該 level 裁切檔幾何逐 level 算
        aff = _std_to_orig_affine(orig_img, t_img, bool(ap_flip))
        if aff is None:
            logger.warning(f'{name} {level}: 原始 CT direction 非純對角，'
                           f'跳過')
            continue
        A, b = aff
        # Resample 的 transform 方向是 dest(原始) physical -> src(std)
        # physical：P_std = A·(P_orig - b) = A·P_orig - A*b（A 對角 ±1）
        tr = _affine_sitk(A, -A * b)
        # 重編號：Lx VBODY=2x-1、SP=2x（x = L1..L6 的 1-based）；
        # other(3) 不納入
        new = np.zeros_like(unrot)
        new[unrot == 1] = 2 * li + 1
        new[unrot == 2] = 2 * li + 2
        maps.append({'A': [float(v) for v in A],
                     'b': [round(float(v), 3) for v in b]})
        if not new.any():
            logger.info(f'{name} {level}: 無 VBODY / SP voxel，不計入')
            continue
        lvl_img = sitk.GetImageFromArray(new)
        lvl_img.CopyInformation(t_img)
        arr = sitk.GetArrayFromImage(_resample_nearest(orig_img, lvl_img, tr))
        merged = np.where(arr > 0, arr, merged)
        info['orig'] = {'vbody': int((arr == 2 * li + 1).sum()),
                        'sp': int((arr == 2 * li + 2).sum())}
        level_info[level] = info
        check = (f'fwd_mismatch={info["forward_mismatch"]:.1e}'
                 if info.get('forward_mismatch') is not None
                 else f'containment={info["containment"]:.4f}')
        logger.info(f'{name} {level}: {check} '
                    f'bone_frac={info["bone_frac"]:.4f} '
                    f'std->orig A={np.round(A, 3)} b={np.round(b, 2)} '
                    f'orig vbody={info["orig"]["vbody"]} '
                    f'sp={info["orig"]["sp"]} (labels {2 * li + 1}/{2 * li + 2})')

    if not level_info:
        return None, 'no usable level', {}, {}
    if int(merged.sum()) == 0:
        return None, 'merged label empty after resample', level_info, {}

    # 對位檢查：label 應落在原始 CT 體內（HU > -100）且主要是骨頭
    #（median HU >= 100；鬆質骨 HU 可 < 100，故不用 bone 比例作硬門檻）
    ct = sitk.GetArrayFromImage(orig_img)
    hu = ct[merged > 0]
    frac_body = float((hu > -100).mean())
    med_hu = float(np.median(hu))
    frac_bone = float((hu > 100).mean())
    vol_info = {'in_body_frac': frac_body, 'bone_frac': frac_bone,
                'median_hu': med_hu, 'std_to_orig': maps, 'ap_flip': bool(ap_flip)}
    if frac_body < 0.98 or med_hu < 100:
        logger.warning(f'{name}: 對位檢查 in_body={frac_body:.4f} '
                       f'median HU={med_hu:.0f}，疑似錯位，跳過')
        return None, f'alignment check failed ' \
                     f'(in_body={frac_body:.3f}, median_hu={med_hu:.0f})', \
                     level_info, vol_info
    logger.info(f'{name}: alignment OK in_body={frac_body:.4f} '
                f'bone_frac={frac_bone:.4f} median HU={med_hu:.0f}')

    out_path = os.path.join(dest_dir, key, f'{name}.nii.gz')
    if os.path.exists(out_path) and not force:
        return out_path, 'exists (skipped)', level_info, vol_info
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    out_img = sitk.GetImageFromArray(merged)
    out_img.CopyInformation(orig_img)
    sitk.WriteImage(out_img, out_path)
    logger.info(f'{name}: saved {out_path} (levels={list(level_info)})')
    return out_path, 'written', level_info, vol_info


def find_original_ct(name):
    """在 DATA_ROOT 子目錄裡找 <name>.nii.gz；回傳 (path, 子目錄) 或
    (None, None)。"""
    if not os.path.isdir(DATA_ROOT):
        return None, None
    for key in sorted(os.listdir(DATA_ROOT)):
        p = os.path.join(DATA_ROOT, key, f'{name}.nii.gz')
        if os.path.isfile(p):
            return p, key
    return None, None


def list_volumes_with_labels(root):
    """root 下有 rotated/*_label.nii.gz 的 volume 目錄（依名稱排序）。"""
    vols = []
    if not os.path.isdir(root):
        return vols
    for d in sorted(os.listdir(root)):
        rot = os.path.join(root, d, 'rotated')
        if os.path.isdir(rot) and any(
                f.endswith('_label.nii.gz') for f in os.listdir(rot)):
            vols.append(os.path.join(root, d))
    return vols


def main():
    parser = argparse.ArgumentParser(
        description='Merge unrotated lumbar VBODY/SP labels into one '
                    'original-CT coordinate .nii.gz per patient.')
    parser.add_argument('--roots', nargs='*', default=DEFAULT_ROOTS,
                        help='Output generation dir(s) to scan (newest first).')
    parser.add_argument('--dest', default=DEFAULT_DEST,
                        help=f'Destination dir (default: {DEFAULT_DEST})')
    parser.add_argument('--force', action='store_true',
                        help='Overwrite existing outputs.')
    parser.add_argument('--min-contain', type=float, default=0.85,
                        help='Fallback per-level check: skip a level whose '
                             'unrotation round-trip containment is below '
                             'this (default: 0.85).')
    parser.add_argument('--name', default=None,
                        help='Process only this volume name (debug).')
    args = parser.parse_args()

    os.makedirs(LOG_DIR, exist_ok=True)
    log_path = os.path.join(LOG_DIR, f'xfr_orig_labels_{time.strftime("%Y%m%d_%H%M%S")}.log')
    setup_tee(log_path)
    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'Roots: {args.roots}')
    logger.info(f'Dest: {args.dest}')
    logger.info('Label scheme: Lx VBODY=2x-1, SP=2x (L1:1/2, L2:3/4, '
                'L3:5/6, L4:7/8, L5:9/10, L6:11/12); other bone excluded; '
                'uint8; geometry = original CT')

    os.makedirs(args.dest, exist_ok=True)
    ap_flip_db = load_ap_flip()
    logger.info(f'ap_flip metadata: {len(ap_flip_db)} entries '
                f'({sum(1 for v in ap_flip_db.values() if v)} flipped)')
    done, summary, no_ap = set(), {}, set()
    summary_path = os.path.join(args.dest, 'summary.json')
    if os.path.exists(summary_path):
        try:
            with open(summary_path) as f:
                summary = json.load(f)
        except (json.JSONDecodeError, OSError):
            summary = {}

    for root in args.roots:
        vols = list_volumes_with_labels(root)
        if args.name is not None:
            vols = [v for v in vols if os.path.basename(v) == args.name]
        logger.info(f'Scanning {root}: {len(vols)} volume(s) with rotated labels')
        for vol_dir in vols:
            name = os.path.basename(vol_dir)
            if name in done:
                logger.info(f'[skip] {name}: 已由較新世代處理')
                continue
            done.add(name)

            orig_path, key = find_original_ct(name)
            if orig_path is None:
                logger.warning(f'[skip] {name}: 找不到原始 CT '
                               f'({DATA_ROOT}*/{name}.nii.gz)')
                continue
            ap_flip = ap_flip_db.get(name)
            if ap_flip is None and name not in no_ap:
                no_ap.add(name)
                logger.warning(f'{name}: metadata db 缺 ap_flip，假設 False '
                               f'（若為 AP 翻轉個案將對位失敗並被抓出）')
                ap_flip = False
            out_path = os.path.join(args.dest, key, f'{name}.nii.gz')
            if os.path.exists(out_path) and not args.force:
                logger.info(f'[skip] {name}: {out_path} 已存在')
                continue

            try:
                orig_img = sitk.ReadImage(orig_path)
                t0 = time.time()
                out, status, level_info, vol_info = process_volume(
                    vol_dir, orig_img, key, args.dest, args.min_iou, args.force,
                    ap_flip)
                dt = time.time() - t0
                common = {'status': status, 'source_root': root,
                          'vol_info': vol_info,
                          'updated': time.strftime('%Y-%m-%d %H:%M:%S')}
                if out is None:
                    logger.warning(f'[skip] {name}: {status} ({dt:.1f}s)')
                    summary[name] = {**common, 'dest_root': DATA_ROOT,
                                     'levels': level_info}
                else:
                    logger.info(f'[done] {name}: {status} in {dt:.1f}s')
                    summary[name] = {**common, 'path': out, 'key': key,
                                     'levels': level_info}
            except Exception as e:
                logger.warning(f'[error] {name}: {e}')
                summary[name] = {'status': f'error: {e}', 'source_root': root,
                                 'updated': time.strftime('%Y-%m-%d %H:%M:%S')}

    with open(summary_path, 'w') as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)
    logger.info('All done.')


if __name__ == '__main__':
    main()