#!/home/xfr/.conda/envs/cbt/bin/python
"""一次性的 layout 遷移：把已產出 volume 的「未旋轉 <level> 檔」移入 crop/ 子資料夾。

新 layout（2026-09-08 起 xfr_preprocess 直接產出，無須再跑本腳本）：
  <vol>/crop/<L>_binary|_smd|_smd_resampled|_binary_sdf|_binary_nn|_roi.nii.gz
  <vol>/crop/<L>_planes.png        （原 <vol>/lumbar/）
  <vol>/rotated/<L>_*.nii.gz       （不變）
  <vol>/transform.json、<name>_labels.txt（頂層，不變）

本腳本供改動前已產出、尚未重新處理的 volume 用（例：standardized-xfr-3）：
  1) <vol>/ 頂層的 <L>_*.nii.gz（L = LABEL_MAP level 名）移到 <vol>/crop/
  2) <vol>/lumbar/*.png 移到 <vol>/crop/，lumbar/ 清空後刪除
Idempotent：已遷就的 volume 重跑無作用；crop/ 已有同名檔時不覆蓋、跳過該檔
（頂層檔保留，人工核對）。rotated/ 與其他頂層檔不動。

讀端（xfr_orig_labels / xfr_inverse_transform / xfr_cbt_native 等）以
imaging.transforms.level_file_path 同時認 crop/ 與頂層，舊世代
（standardized-xfr-2 等）不遷移也能讀；如需一致化可對該 root 跑本腳本。

Usage:
  python xfr_migrate_crop.py [ROOT] [--dry-run]
"""
import argparse
import os
import shutil
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from config.constant import LABEL_MAP

DEFAULT_ROOT = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3'
LEVEL_SET = set(LABEL_MAP.values())


def migrate_volume(vol_dir, dry_run=False):
    """回傳 (moved, kept) 兩個清單（人讀的說明文字）。"""
    crop_dir = os.path.join(vol_dir, 'crop')
    moved, kept = [], []

    for f in sorted(os.listdir(vol_dir)):
        if not f.endswith('.nii.gz'):
            continue
        if f.split('_', 1)[0] not in LEVEL_SET:
            kept.append(f)
            continue
        dst = os.path.join(crop_dir, f)
        if os.path.exists(dst):
            kept.append(f'{f}（crop/ 已有同名檔，未動）')
            continue
        src = os.path.join(vol_dir, f)
        if dry_run:
            moved.append(f'{f} (dry-run)')
        else:
            os.makedirs(crop_dir, exist_ok=True)
            shutil.move(src, dst)
            moved.append(f)

    lumbar_dir = os.path.join(vol_dir, 'lumbar')
    if os.path.isdir(lumbar_dir):
        for f in sorted(os.listdir(lumbar_dir)):
            dst = os.path.join(crop_dir, f)
            if os.path.exists(dst):
                kept.append(f'lumbar/{f}（crop/ 已有同名檔，未動）')
                continue
            if dry_run:
                moved.append(f'lumbar/{f} (dry-run)')
            else:
                os.makedirs(crop_dir, exist_ok=True)
                shutil.move(os.path.join(lumbar_dir, f), dst)
                moved.append(f'lumbar/{f}')
        if not dry_run:
            try:
                os.rmdir(lumbar_dir)
            except OSError:
                kept.append('lumbar/（仍有其他檔，未刪）')
    return moved, kept


def main():
    parser = argparse.ArgumentParser(
        description='Migrate per-level pre-rotated files from <vol>/ top level '
                    'to <vol>/crop/ (plus lumbar/*.png -> crop/).')
    parser.add_argument('root', nargs='?', default=DEFAULT_ROOT,
                        help=f'standardized generation dir (default: {DEFAULT_ROOT})')
    parser.add_argument('--dry-run', action='store_true',
                        help='Only report what would be moved.')
    args = parser.parse_args()

    if not os.path.isdir(args.root):
        sys.exit(f'not a directory: {args.root}')
    vols = sorted(d for d in os.listdir(args.root)
                  if os.path.isdir(os.path.join(args.root, d)))
    tag = '[dry-run] ' if args.dry_run else ''
    n_done = n_empty = 0
    for i, vol in enumerate(vols, 1):
        vol_dir = os.path.join(args.root, vol)
        moved, kept = migrate_volume(vol_dir, dry_run=args.dry_run)
        if not moved:
            n_empty += 1
            print(f'{tag}[{i}/{len(vols)}] -     {vol}  (no file to move)')
            continue
        n_done += 1
        print(f'{tag}[{i}/{len(vols)}] ok    {vol}  moved {len(moved)} file(s)')
        for m in moved:
            print(f'    -> {m}')
        for k in kept:
            print(f'    !! kept: {k}')
    print(f'\n{tag}done: {n_done} volume(s) migrated, {n_empty} untouched '
          f'out of {len(vols)}')


if __name__ == '__main__':
    main()