#!/home/xfr/.conda/envs/cbt/bin/python
"""補記既有 volume 的 transform.json rotated 段平面（sym / endplate）。

core/optimizer 的 xfr 管線（run_pso_torch_xfr / run_de_torch_xfr）不再於影像
上重新估算兩平面（best_symmetry_plane / best_upper_endplate_plane），改直接
讀 <vol>/transform.json levels[L*]['rotated']['sym' / 'endplate']
（imaging.transforms.load_rotated_planes_from_path）：
  - 新世代：xfr_preprocess._write_rotated_level 旋轉時即記錄（rotated/ 局部系）；
  - 舊世代（進度 progress.json 已標完成、重跑 preprocess 會整支跳過）：跑本腳本
    補記——用與舊 optimizer 相同函式、相同預設值，對
    rotated/<level>_binary_sdf.nii.gz（完整 mask、SP 移除前）重新估算，
    併入既有 rotated 段（只加 sym / endplate 兩鍵，不動 template / R / center /
    start / size 等既有鍵、不重寫任何 nii）。補記後 optimizer 讀到的平面與
    改前行為一致（同輸入、同演算法）。

預設跳過已有 sym 記錄的 level（部分失敗後重跑只補缺者）；--force 全數重記。

Usage:
  python xfr_record_planes.py [--force] [--dir <standardized dir>] [volume_id ...]
  volume_id 可用完整 ID 或末段（如 0001）；省略 = 全部 volume。
"""
import argparse
import logging
import os
import re
import sys
import time

import numpy as np
import SimpleITK as sitk

from imaging.orientation import best_symmetry_plane, best_upper_endplate_plane
from imaging.transforms import load_transform, save_transform


STANDARDIZED_DIR = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3/'

LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs')
logger = logging.getLogger('xfr_record_planes')

# 基準系健全閾值：對齊良好時 theta_v / tau_y 應近似 0（xfr_preprocess 的 R
# 已把鏡稱面法線轉到 +x、終板法線轉到 z>0）
SANITY_DEG = 5.0


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):
    log_fh = open(log_path, 'a', buffering=1)
    sys.stdout = _Tee(sys.stdout, log_fh)
    sys.stderr = _Tee(sys.stderr, log_fh)


def rotated_levels(vol_dir):
    """該 volume rotated/ 中有 _binary_sdf.nii.gz 的 lumbar level（排序）。"""
    rotated_dir = os.path.join(vol_dir, 'rotated')
    if not os.path.isdir(rotated_dir):
        return ()
    levels = set()
    for fn in os.listdir(rotated_dir):
        m = re.fullmatch(r'(L[1-9]\d*)_binary_sdf\.nii\.gz', fn)
        if m:
            levels.add(m.group(1))
    return tuple(sorted(levels, key=lambda lv: int(lv[1:])))


def record_level(vol_dir, level, sdf_path):
    """估算該 level 兩平面併入 transform.json rotated 段；回傳 (sym, endplate)。"""
    arr = sitk.GetArrayFromImage(sitk.ReadImage(sdf_path))
    sym = best_symmetry_plane(arr)
    endplate = best_upper_endplate_plane(arr)

    meta = load_transform(vol_dir)
    lv = meta['levels'].setdefault(level, {})
    rot = lv.setdefault('rotated', {})
    rot['sym'] = sym
    rot['endplate'] = endplate
    save_transform(vol_dir, meta)

    s_nx, s_ny, _ = sym['normal']
    theta_v = float(np.degrees(np.arctan2(s_ny, s_nx)))
    if endplate is not None:
        e_nx, e_ny, e_nz = endplate['normal']
        tau_y = float(np.degrees(np.arctan2(e_ny, e_nz)))
    else:
        tau_y = None
    note = f'endplate={endplate is not None}'
    if abs(theta_v) > SANITY_DEG or (tau_y is not None and abs(tau_y) > SANITY_DEG):
        logger.warning(f'{os.path.basename(vol_dir)} {level}: 基準系殘差偏大 '
                       f'(theta_v={theta_v:+.2f} deg, tau_y='
                       f'{tau_y if tau_y is None else round(tau_y, 2)}) — '
                       f'該 level 的 xfr_preprocess 對齊可能失效，建議人工核對')
    else:
        logger.info(f'{os.path.basename(vol_dir)} {level}: recorded '
                    f'(mirror ratio={sym["ratio"]:.3f}, theta_v={theta_v:+.2f} deg, '
                    f'tau_y={tau_y if tau_y is None else round(tau_y, 2)}, {note})')
    return sym, endplate


def main():
    parser = argparse.ArgumentParser(
        description='Record sym/endplate planes into transform.json for '
                    'existing volumes (old generation, before xfr_preprocess '
                    'recording them at rotation time).')
    parser.add_argument('volume_ids', nargs='*',
                        help='Full ID or trailing segment (e.g. 0001); '
                             'default: all volumes')
    parser.add_argument('--dir', default=STANDARDIZED_DIR,
                        help='Standardized output dir (default: standardized-xfr-3)')
    parser.add_argument('--force', action='store_true',
                        help='Re-record levels that already have sym/endplate')
    args = parser.parse_args()

    os.makedirs(LOG_DIR, exist_ok=True)
    log_path = os.path.join(
        LOG_DIR, f'xfr_record_planes_{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)}')

    if not os.path.isdir(args.dir):
        sys.exit(f'Standardized dir not found: {args.dir}')
    volumes = [d for d in sorted(os.listdir(args.dir))
               if os.path.isdir(os.path.join(args.dir, d))]
    if args.volume_ids:
        keys = [v.lower() for v in args.volume_ids]
        vols = [v for v in volumes
                if v.lower() in keys or v.rsplit('.', 1)[-1].lower() in keys]
        if not vols:
            sys.exit(f'No volume matches: {", ".join(args.volume_ids)}')
        volumes = vols
    logger.info(f'{len(volumes)} volume(s) under {args.dir}'
                + (' (force)' if args.force else ''))

    t0 = time.time()
    n_rec, n_skip_done, n_err, n_no_levels = 0, 0, 0, 0
    no_ep = []
    for vid in volumes:
        vol_dir = os.path.join(args.dir, vid)
        levels = rotated_levels(vol_dir)
        if not levels:
            n_no_levels += 1
            continue
        try:
            meta = load_transform(vol_dir)
        except FileNotFoundError:
            logger.warning(f'{vid}: transform.json 缺失，跳過 '
                           f'（舊世代可先在 xfr_inverse_transform.py 加 --rebuild）')
            n_err += 1
            continue
        for level in levels:
            rot = meta['levels'].get(level, {}).get('rotated') or {}
            if not args.force and rot.get('sym') is not None:
                n_skip_done += 1
                continue
            sdf_path = os.path.join(vol_dir, 'rotated', f'{level}_binary_sdf.nii.gz')
            try:
                _, ep = record_level(vol_dir, level, sdf_path)
                n_rec += 1
                if ep is None:
                    no_ep.append(f'{vid} {level}')
            except Exception as e:
                n_err += 1
                logger.error(f'{vid} {level}: {e}')
    dt = time.time() - t0
    logger.info(f'Done in {dt / 60:.1f} min: recorded {n_rec} level(s), '
                f'skipped(already) {n_skip_done}, no rotated levels '
                f'{n_no_levels} volume(s), errors {n_err}')
    if no_ep:
        logger.warning(f'{len(no_ep)} level(s) 無上終板面記錄（資料不足），'
                       f'optimizer 對應 level 終板入口剪除停用: {", ".join(no_ep[:10])}'
                       f'{" ..." if len(no_ep) > 10 else ""}')


if __name__ == '__main__':
    main()