#!/home/xfr/.conda/envs/cbt/bin/python
"""把 PSO 找到的 CBT 螺絲位置（旋轉 0.5mm 標準系）映回原始
（未旋轉、未重取樣）native CT 空間，整支 volume 各 level 的螺絲
（L1-L5 x L/R，最多 10 支）存成單一 label 體積：

    Output_dir/<run_date>/<volume_id>/cbt.nii.gz

label 值：L1L=1 L1R=2 L2L=3 L2R=4 L3L=5 L3R=6 L4L=7 L4R=8 L5L=9 L5R=10
（0 = 背景）。

座標鏈（純 index 空間；不依賴 standardized 輸出的物理 header ——
standardize_affine 的 origin 處理不可靠，2026-09-08 已驗證）：
  rotated disk r（(x,y,z) index）
    -> template 0.5mm:   t = R^T (r + fstart - c) + c   [行向量: (r+fstart-c) @ R + c]
    -> 記憶體 0.5mm 全域: g0 = bbox2s + (wx-1-tx, wy-1-ty, tz)
    -> ap_flip 時:        g0y = N05y - 1 - g0y
    -> native index:      rint(g0 * 0.5 / sn)

其中每 level 的 R/c 由 template 骨頭 mask（smd_resampled < 0.5，與
_write_rotated_level 的輸入同定義；退化時退回 _binary_nn）重算；
fstart = rotated 檔 origin 反映到 template index 的整數；
bbox2 = native label 線性重取樣 0.5mm（>0.5）的最大 26-連通區域 bbox，
與 seg_bone 同定義（template/roi/binary_sdf 都裁在這個 bbox2 上）。

螺絲參數化（rotated 系、0.5mm index）：pos = [z, y, x, az°, alt°, d mm, L mm]；
方向 d_v = (cos az sin alt, sin az sin alt, cos alt)；末端 = p0 + L/0.5 * d_v；
柱半徑 = d/2 mm。

驗證（2026-09-08）：整顆骨頭 15/15 (volume, level) 100% 落在 native label
（/tmp/kilo/validate_final.py）；螺絲柱體全部點 20/20 level-side
in-label >= 96%（2-voxel 膨脹）（/tmp/kilo/validate_screws2.py）。
"""

import argparse
import json
import logging
import os
import sys

import numpy as np
import SimpleITK as sitk
from tinydb import TinyDB, Query

from config.constant import LABEL_MAP
from imaging.resample import resample_img
from imaging.segmentation import _largest_cc_bbox
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

standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3'
data_root = '/mnt/1220/Public/dataset/Spine/CTSpine1K/data/'
label_root = '/mnt/1220/Public/dataset/Spine/CTSpine1K/label/'
SUBDIRS = {
    'colon': 'conlon',
    'COVID-19': 'COVID-19',
    'HNSCC-3DCT-RT_neck': 'HNSCC-3DCT-RT_neck',
    'liver': 'Liver',
}
Output_dir = '/mnt/1248/open/cyrou/Output'

_PROJ_DIR = os.path.dirname(os.path.abspath(__file__))
LOG_DIR = os.path.join(_PROJ_DIR, 'logs')
SIDE_RESULT_DIR = os.path.join(LOG_DIR, 'side_results')
META_DB = os.path.join(_PROJ_DIR, 'xfr_image_metadata.json')

LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5')
LEVEL_LABEL_VAL = {v: int(k) for k, v in LABEL_MAP.items() if v in LEVELS}  # {'L1': 20, ...}

logger = logging.getLogger('xfr_cbt_native')


def find_native_paths(volume_id):
    """回傳 (ct_path, label_path, ap_flip)；找不到 native 檔對時 raise。"""
    ap_flip = False
    try:
        meta = TinyDB(META_DB, access_mode='r').table('images') \
            .get(Query().name == volume_id)
        if meta is None:
            logger.warning(f'{volume_id}: not in metadata db; ap_flip defaults to False')
        else:
            ap_flip = bool(meta.get('ap_flip', False))
    except Exception as e:
        logger.warning(f'{volume_id}: metadata db unavailable ({e}); ap_flip defaults to False')
    for sub, lab_sub in SUBDIRS.items():
        ct = f'{data_root}{sub}/{volume_id}.nii.gz'
        lb = f'{label_root}{lab_sub}/{volume_id}_seg.nii.gz'
        if os.path.isfile(ct) and os.path.isfile(lb):
            return ct, lb, ap_flip
    raise FileNotFoundError(f'{volume_id}: no native CT/label pair under {data_root}')


def level_geometry(volume_id, level, lb_img, lb_arr):
    """每 level 的 (R, c, fstart, bbox2s, wx, wy)；與預處理/驗證腳本同定義。"""
    vol_dir = os.path.join(standardized_dir, volume_id)
    # 未旋轉檔：新世代在 crop/ 子資料夾、舊世代在頂層
    smd_img = sitk.ReadImage(level_file_path(vol_dir, level, 'smd_resampled'))
    smd = sitk.GetArrayFromImage(smd_img).astype(np.float32)
    if not (np.isfinite(smd).all() and smd.min() < 0 and smd.max() > 0):
        b = (sitk.GetArrayFromImage(
            sitk.ReadImage(level_file_path(vol_dir, level, 'binary_nn'))) > 0).astype(np.uint8)
    else:
        b = (smd < 0.5).astype(np.uint8)
    sym = best_symmetry_plane(b)
    symp = best_upper_endplate_plane(b)
    R, c = compute_normalizing_rotation(b, sym, symp)
    rot_img = sitk.ReadImage(f'{vol_dir}/rotated/{level}_binary_sdf.nii.gz')
    fstart = np.round(np.array(
        smd_img.TransformPhysicalPointToIndex(rot_img.GetOrigin()))).astype(float)

    lv = LEVEL_LABEL_VAL[level]
    bin_img = sitk.GetImageFromArray((lb_arr == lv).astype(np.uint8))
    bin_img.CopyInformation(lb_img)
    bin_lin = resample_img(sitk.Cast(bin_img, sitk.sitkFloat32))
    m_full = sitk.GetImageFromArray((sitk.GetArrayFromImage(bin_lin) > 0.5).astype(np.uint8))
    m_full.CopyInformation(bin_lin)
    cc = _largest_cc_bbox(m_full)
    if cc is None:
        raise ValueError(f'{volume_id} {level}: empty native level mask (label {lv})')
    _, bbox2 = cc
    wx, wy = smd_img.GetSize()[0], smd_img.GetSize()[1]
    return R, c, fstart, np.array(bbox2[:3], float), wx, wy


def rotated_to_native(r_xyz, geom, N05, sn, ap_flip):
    """r_xyz: (N,3) rotated disk (x,y,z) index -> (N,3) native 連續 index。"""
    R, c, fstart, bbox2s, wx, wy = geom
    t = (np.asarray(r_xyz, float) + fstart - c) @ R + c
    g0 = np.empty_like(t)
    g0[:, 0] = bbox2s[0] + (wx - 1 - t[:, 0])
    g0[:, 1] = bbox2s[1] + (wy - 1 - t[:, 1])
    g0[:, 2] = bbox2s[2] + t[:, 2]
    if ap_flip:
        g0[:, 1] = N05[1] - 1 - g0[:, 1]
    return g0 * (0.5 / sn)


def screw_voxel_xyz(pos):
    """pos = [z, y, x, az°, alt°, d, L]（rotated 系 0.5mm index）。
    回傳柱體內 (x,y,z) float voxel 座標（N,3）。"""
    z, y, x, az, alt, d, L = (float(v) for v in pos[:7])
    azr, altr = np.radians(az), np.radians(alt)
    ca, sa = np.cos(azr), np.sin(azr)
    ct_, st_ = np.cos(altr), np.sin(altr)
    dv = np.array([ca * st_, sa * st_, ct_])
    e1 = np.array([ca * ct_, sa * ct_, -st_])
    e2 = np.array([-sa, ca, 0.0])
    p0 = np.array([x, y, z])
    p1 = p0 + (L / 0.5) * dv
    rad = d / 0.5 + 3.0
    lo = np.floor(np.minimum(p0, p1)) - rad
    hi = np.ceil(np.maximum(p0, p1)) + rad
    xs, ys, zs = np.meshgrid(np.arange(lo[0], hi[0] + 1),
                             np.arange(lo[1], hi[1] + 1),
                             np.arange(lo[2], hi[2] + 1), indexing='xy')
    P = np.stack([xs.ravel(), ys.ravel(), zs.ravel()], 1)
    dP = P - p0
    xr, yr, zr = dP @ e1, dP @ e2, dP @ dv
    rrad = d  # (d/2 mm 半徑) / 0.5mm 每 voxel = d 個 voxel
    m = (xr ** 2 + yr ** 2 <= rrad ** 2) & (zr >= 0) & (zr <= L / 0.5)
    return P[m]


def write_volume_cbt(volume_id, run_id, output_root=Output_dir, date=None):
    """把一個 volume 的 side_results 螺絲全部映回 native 空間寫 cbt.nii.gz。
    回傳 (path, n_screws)；該 run 無此 volume 時回傳 (None, 0)。"""
    date = date or run_id[:8]
    side_vol_dir = os.path.join(SIDE_RESULT_DIR, run_id, volume_id)
    if not os.path.isdir(side_vol_dir):
        logger.warning(f'{volume_id}: no side results under {side_vol_dir}; cbt skipped')
        return None, 0
    if not any(os.path.isfile(os.path.join(side_vol_dir, f'{l}_{s}.json'))
               for l in LEVELS for s in ('L', 'R')):
        logger.warning(f'{volume_id}: no <level>_<side>.json in {side_vol_dir}; cbt skipped')
        return None, 0

    ct_path, lb_path, ap_flip = find_native_paths(volume_id)
    ct = sitk.ReadImage(ct_path)
    lb_img = sitk.ReadImage(lb_path)
    lb_arr = sitk.GetArrayFromImage(lb_img)
    Nn = np.array(ct.GetSize(), float)
    sn = np.array(ct.GetSpacing(), float)
    N05 = np.maximum(1, np.ceil(Nn * sn / 0.5 - 1e-6).astype(int))

    out_arr = np.zeros(lb_arr.shape, np.uint8)
    geom_cache = {}
    n_screws, skipped = 0, []
    for li, level in enumerate(LEVELS):
        for side in ('L', 'R'):
            jp = os.path.join(side_vol_dir, f'{level}_{side}.json')
            if not os.path.isfile(jp):
                continue
            pos = json.load(open(jp))['position']
            try:
                if level not in geom_cache:
                    geom_cache[level] = level_geometry(volume_id, level, lb_img, lb_arr)
                geom = geom_cache[level]
                cyl = screw_voxel_xyz(pos)
                nat = rotated_to_native(cyl, geom, N05, sn, ap_flip)
                ni = np.rint(nat).astype(int)
                valid = (ni >= 0).all(1) & (ni < np.array(lb_arr.shape[::-1])).all(1)
                idx = ni[valid]
                val = li * 2 + (1 if side == 'R' else 0)
                out_arr[idx[:, 2], idx[:, 1], idx[:, 0]] = val
                n_screws += 1
            except Exception as e:
                logger.error(f'{volume_id} {level}_{side}: {e}')
                skipped.append(f'{level}_{side} ({e})')

    if n_screws == 0:
        raise ValueError(f'{volume_id}: no screw could be written (skipped: {skipped or "none read"})')
    out_dir = os.path.join(output_root, date, volume_id)
    os.makedirs(out_dir, exist_ok=True)
    out_img = sitk.GetImageFromArray(out_arr)
    out_img.CopyInformation(ct)
    out_path = os.path.join(out_dir, 'cbt.nii.gz')
    sitk.WriteImage(out_img, out_path)
    if skipped:
        logger.warning(f'{volume_id}: failed sides: {", ".join(skipped)}')
    logger.info(f'{volume_id}: {n_screws}/10 screws -> {out_path} (ap_flip={ap_flip})')
    return out_path, n_screws


def main():
    parser = argparse.ArgumentParser(
        description='Write <Output_dir>/<date>/<volume_id>/cbt.nii.gz (screws mapped to native space) '
                    'from a side_results run')
    parser.add_argument('run_id', help='e.g. 20260907_230652')
    parser.add_argument('volumes', nargs='*',
                        help='volume ids (full or trailing digits); default: all in the run')
    parser.add_argument('--date', default=None, help='output date dir (default: run_id[:8])')
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO,
                        format='%(asctime)s [%(levelname)s] %(name)s: %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S')
    run_dir = os.path.join(SIDE_RESULT_DIR, args.run_id)
    if not os.path.isdir(run_dir):
        sys.exit(f'side results dir not found: {run_dir}')
    run_vols = sorted(d for d in os.listdir(run_dir) if os.path.isdir(os.path.join(run_dir, d)))
    if args.volumes:
        vols = [v for v in run_vols
                if v in args.volumes or v.rsplit('.', 1)[-1] in args.volumes]
        if not vols:
            sys.exit(f'none of {args.volumes} found in run {args.run_id}')
    else:
        vols = run_vols
    ok, fail = 0, 0
    for vid in vols:
        try:
            write_volume_cbt(vid, args.run_id, date=args.date)
            ok += 1
        except Exception as e:
            fail += 1
            logger.error(f'{vid}: {e}')
    print(f'done: {ok} written, {fail} failed / {len(vols)} volume(s)')


if __name__ == '__main__':
    main()