#!/home/xfr/.conda/envs/cbt/bin/python
"""掃 standardize 輸出目錄中「前後（AP）方向翻轉」的個案（prone 伏位掃描，
前側落在 y 小側，例：colon 0003），供重跑修正用。

判定：對每個 volume 的 L1~L6 輸出遮罩（優 _binary_sdf、次 _binary_nn、
再 _binary）個別跑 orientation.anterior_y_side（中線帶椎管兩側質量比較，
見該函式 docstring），多 level 票決：
  flip    : y_min 票 > y_max 票（前後翻轉，需重跑）
  ok      : y_max 票 > y_min 票（方向正常）
  mixed   : 平手（需人工確認）
  unknown : 全部無法判定（無明顯前後質量差，如鏡稱面異常案例）

Usage:
  python xfr_reprocess_ap.py <output_dir>            # 只回報
  python xfr_reprocess_ap.py <output_dir> --fix      # 另刪 flip volume 的輸出
                                                     # 資料夾 + progress.json
                                                     # 條目，之後重跑
                                                     # xfr_preprocess.py 即可
"""
import argparse
import json
import os
import shutil
import sys
import time

import SimpleITK as sitk

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from imaging.orientation import anterior_y_side
from imaging.transforms import level_file_path

LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5', 'L6')
MASK_SOURCES = ('binary_sdf', 'binary_nn', 'binary')


def volume_decision(vol_dir):
    """回傳 (decision, per_level dict)。decision ∈ flip/ok/mixed/unknown/nomask。"""
    per = {}
    for lvl in LEVELS:
        # 未旋轉遮罩：新世代在 crop/ 子資料夾、舊世代在頂層
        for src in MASK_SOURCES:
            p = level_file_path(vol_dir, lvl, src)
            if os.path.exists(p):
                m = sitk.GetArrayFromImage(sitk.ReadImage(p, sitk.sitkUInt8))
                per[lvl] = anterior_y_side(m)
                break
    if not per:
        return 'nomask', per
    votes = [v for v in per.values() if v is not None]
    if not votes:
        return 'unknown', per
    n_min = votes.count('y_min')
    n_max = votes.count('y_max')
    if n_min > n_max:
        return 'flip', per
    if n_max > n_min:
        return 'ok', per
    return 'mixed', per


def main():
    parser = argparse.ArgumentParser(
        description='Find AP-flipped (prone) volumes in a standardized output dir.')
    parser.add_argument('output_dir')
    parser.add_argument('--fix', action='store_true',
                        help='Also delete flipped volumes\' output dirs and '
                             'their progress.json entries')
    args = parser.parse_args()

    outdir = args.output_dir
    if not os.path.isdir(outdir):
        print(f'not a directory: {outdir}')
        return

    vols = sorted(d for d in os.listdir(outdir)
                  if os.path.isdir(os.path.join(outdir, d)))
    flip, mixed, unknown, ok, nomask = [], [], [], [], []
    t0 = time.time()
    for i, vol in enumerate(vols, 1):
        dec, per = volume_decision(os.path.join(outdir, vol))
        tag = {'flip': 'FLIP', 'ok': 'ok  ', 'mixed': 'MIXED',
               'unknown': '?!?', 'nomask': '-    '}[dec]
        detail = ' '.join(f'{k}={v}' for k, v in per.items())
        print(f'[{i}/{len(vols)}] {tag} {vol}  {detail}')
        {'flip': flip, 'mixed': mixed, 'unknown': unknown,
         'ok': ok, 'nomask': nomask}[dec].append(vol)
        if (i % 25) == 0:
            print(f'  ... {i}/{len(vols)}  ({(time.time()-t0)/60:.1f} min)', flush=True)

    print(f'\n=== Summary: {len(vols)} volumes ===')
    print(f'  ok (normal)        : {len(ok)}')
    print(f'  FLIP (AP-flipped)  : {len(flip)}')
    for v in flip:
        print(f'    - {v}')
    print(f'  mixed (need check) : {len(mixed)}')
    for v in mixed:
        print(f'    - {v}')
    print(f'  unknown (no vote)  : {len(unknown)}')
    for v in unknown:
        print(f'    - {v}')
    print(f'  no level mask      : {len(nomask)}')

    if args.fix and flip:
        # 1) progress.json：刪掉 flip 的條目（備份）
        prog_path = os.path.join(outdir, 'progress.json')
        if os.path.exists(prog_path):
            with open(prog_path) as f:
                prog = json.load(f)
            removed = [v for v in flip if prog.pop(v, None) is not None]
            bak = f'{prog_path}.apfix-{time.strftime("%Y%m%d_%H%M%S")}'
            shutil.copyfile(prog_path, bak)
            with open(prog_path, 'w') as f:
                json.dump(prog, f, indent=2)
            print(f'\nprogress.json: removed {len(removed)} entr(y/ies) '
                  f'[{", ".join(v.split(".")[-1] or v for v in removed)}]; '
                  f'backup {bak}')
        # 2) 刪輸出資料夾
        for v in flip:
            shutil.rmtree(os.path.join(outdir, v))
            print(f'removed {os.path.join(outdir, v)}')
        print('\nNext: rerun `python xfr_preprocess.py` — only the removed '
              'volumes will be reprocessed (with the AP flip applied).')
    elif not args.fix and flip:
        print('\nRerun with --fix to delete the flipped outputs and progress '
              'entries, then run `python xfr_preprocess.py`.')


if __name__ == '__main__':
    main()