#!/home/xfr/.conda/envs/cbt/bin/python
"""
為每個 (volume, level) 的骨頭遮罩（{level}_binary_sdf.nii.gz，缺則 _binary）繪製
X-ray 四視角圖（不畫螺絲），存到 Output/{date}/{volume}/。

輸出檔名：{volume} {level}_CBT.png
（例：1.3.6.1.4.1.9328.50.4.0005 L1_CBT.png）

Usage:
  python xfr_plot_level.py                     # 全部 volume 的 L1~L5
  python xfr_plot_level.py 0005                # 該 volume 的 L1~L5
  python xfr_plot_level.py 0005 L1             # 單一 (volume, level)
  python xfr_plot_level.py --dir <standardized_dir> --output <output_base>
"""

import argparse
import os
import sys

from imaging.transforms import level_file_path
from visualization.res_bone_figure import render_bone_figure

standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr/'
output_base = '/mnt/1248/open2/cyrou/Output'

LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5')

USAGE = 'Usage: python xfr_plot_level.py [volume_id] [level] [--dir D] [--output O]'


def parse_args(argv):
    parser = argparse.ArgumentParser(description='繪製各 lumbar level 的骨頭 X-ray 圖（不畫螺絲）')
    parser.add_argument('volume_id', nargs='?', default=None,
                        help='volume ID（完整 UID 或末段，如 0005）；省略＝全部')
    parser.add_argument('level', nargs='?', default=None,
                        help=f'level（{" / ".join(LEVELS)}）；省略＝全部')
    parser.add_argument('--dir', default=standardized_dir,
                        help=f'standardized 資料夾（預設 {standardized_dir}）')
    parser.add_argument('--output', default=output_base,
                        help=f'輸出根目錄（預設 {output_base}）')
    return parser.parse_args(argv)


def main():
    args = parse_args(sys.argv[1:])

    if args.level is not None:
        level_key = args.level.upper()
        if level_key not in LEVELS:
            print(f'Invalid level: {args.level} (choose from {"/".join(LEVELS)})')
            sys.exit(1)
    if args.level is not None and args.volume_id is None:
        print(f'{USAGE}\nlevel requires volume_id')
        sys.exit(1)

    volumes = [d for d in sorted(os.listdir(args.dir))
               if os.path.isdir(os.path.join(args.dir, d))]
    if args.volume_id is not None:
        key = args.volume_id.lower()
        vols = [v for v in volumes
                if v.lower() == key or v.rsplit('.', 1)[-1] == key]
        if not vols:
            print(f'Volume not found: {args.volume_id}')
            sys.exit(1)
        volumes = vols

    levels = (args.level.upper(),) if args.level else LEVELS
    tasks = [(vid, lvl) for vid in volumes for lvl in levels]
    print(f'{len(volumes)} volume(s) x {len(levels)} level(s) = {len(tasks)} figure(s)',
          flush=True)

    ok, skip = [], []
    for i, (vid, lvl) in enumerate(tasks, 1):
        vol_dir = os.path.join(args.dir, vid)
        # _binary.nii.gz 現為原解析度；0.5mm 用 _binary_sdf（SDF 平滑遮罩）
        # 未旋轉檔：新世代在 crop/ 子資料夾、舊世代在頂層
        sdf_path = level_file_path(vol_dir, lvl, 'binary_sdf')
        binary_path = sdf_path if os.path.exists(sdf_path) \
            else level_file_path(vol_dir, lvl, 'binary')
        cortical_path = level_file_path(vol_dir, lvl, 'cortical')
        path = render_bone_figure(vid, lvl, binary_path, cortical_path,
                                  base_folder=args.output)
        if path is None:
            skip.append((vid, lvl))
            print(f'[{i}/{len(tasks)}] {vid} {lvl}: skipped', flush=True)
        else:
            ok.append((vid, lvl))
            print(f'[{i}/{len(tasks)}] {vid} {lvl}: saved {path}', flush=True)

    print('=' * 60)
    print(f'Done. saved={len(ok)} skipped={len(skip)}')
    if skip:
        print('Skipped (missing/empty mask):')
        for vid, lvl in skip:
            print(f'  {vid} {lvl}')


if __name__ == '__main__':
    main()