CBT_project/core/optimizer.py

import json
import os
import time
from datetime import datetime
import SimpleITK as sitk
import torch
from imaging.orientation import (azimuth_rotation, analyze_vertebral_tilt_contour,
                                  best_symmetry_plane, best_upper_endplate_plane,
                                  segment_spinous_process, segment_vertebral_body)
from config.constant import ALLOWED_DIAMETERS, ALLOWED_LENGTHS
from core.objective import OptimizationContext, make_objective_function, make_objective_function_xfr
from pyswarm import pso
from core.cylinder import generate_cylinder_n_torch, generate_cylinder_o_torch, snap_to_discrete_values, create_coordinate_grid, snap_to_discrete_values_xfr, generate_cylinder_butt_torch
from core.intersection import center_line_intersections_torch
from core.scoring import compute_overlap_ratio_from_cylinder_mask, is_solution_ok, cl_score_torch_xfr
from config.constant import OVERLAP_THRESH
from visualization.res_bone_figure import render_bone_figure

LATERAL_REFINE_MIN_IN_BONE = 0.97

# 入口柱向前（+y，椎體方向）幾個 voxel 內進入 VBODY 即視為「入口在椎體上」：
# 涵蓋椎體後側皮質邊緣（mask 外 1~2 voxel 的分割界線帶）；
# 真正的後側要素與椎體間隔（椎間孔）大於此值，不受影響。
VBODY_ENTRY_EDGE = 4


def _makedirs_retry(path, retries=5, delay=0.5):
    """NFS 上建目錄重試（同 utils.helpers.retry_robust 處理的瞬時錯誤）"""
    for i in range(retries):
        try:
            os.makedirs(path, exist_ok=True)
            return
        except OSError:
            if i == retries - 1:
                raise
            time.sleep(delay)


def refine_lateral_longer(
    z, x, azimuth, altitude, diameter_raw, length_raw,
    side, az_bounds, x_bounds, y_indices,
    image_shape, spacing, device, grid,
    cortical_tensor, spine_tensor,
    vbody_tensor=None,
):
    """
    Deterministic local refinement after PSO: try aiming more laterally and
    using a longer screw. Only accepts a candidate if it stays
    >= LATERAL_REFINE_MIN_IN_BONE inside bone AND improves the score.

    left  (x-lower half): more lateral = larger azimuth, entry shifted toward -x
    right (x-upper half): more lateral = smaller azimuth, entry shifted toward +x
    """
    def _score_candidate(cand):
        z_c, x_c, az_c, alt_c, d_c, L_c = cand
        y_c = y_indices[round(z_c), round(x_c)]
        if y_c < 0:
            return None
        cyl = generate_cylinder_n_torch(d_c, L_c, z_c, y_c, x_c, az_c, alt_c,
                                        image_shape, spacing, device, grid)
        cyl_o = generate_cylinder_o_torch(d_c, L_c, z_c, y_c, x_c, az_c, alt_c,
                                          image_shape, spacing, device, grid)
        if cyl.sum().item() == 0:
            return None
        inter, _ = center_line_intersections_torch(z_c, y_c, x_c, az_c, alt_c,
                                                   L_c, spine_tensor, spacing, device)
        cyl_butt = generate_cylinder_butt_torch(d_c, z_c, y_c, x_c, az_c, alt_c,
                                                image_shape, spacing, device, grid)
        loss = cl_score_torch_xfr(cortical_tensor, spine_tensor, cyl, cyl_o, inter,
                                  diameter=d_c,
                                  vbody_tensor=vbody_tensor,
                                  cylinder_butt_torch=cyl_butt)
        in_bone = ((spine_tensor == 1) & (cyl == 1)).sum().item() / cyl.sum().item()
        return {'pos': cand, 'loss': loss, 'in_bone': in_bone}

    d_snap, L_snap = snap_to_discrete_values_xfr(diameter_raw, length_raw)
    az_lo, az_hi = az_bounds
    x_lo, x_hi = x_bounds

    az_steps = [0.0, 5.0, 10.0, 15.0] if side == "L" else [0.0, -5.0, -10.0, -15.0]
    x_shifts = [0.0, -4.0, -8.0] if side == "L" else [0.0, 4.0, 8.0]

    best = _score_candidate((z, x, azimuth, altitude, d_snap, L_snap))
    if best is None:
        return z, x, azimuth, altitude, diameter_raw, length_raw, False, None

    for ds in az_steps:
        az_c = min(az_hi - 0.01, max(az_lo + 0.01, azimuth + ds))
        if abs(az_c - azimuth) < 0.5 and ds != 0:
            continue
        for dx in x_shifts:
            x_c = min(x_hi, max(x_lo, x + dx))
            if abs(x_c - x) < 0.5 and dx != 0:
                continue
            for L_c in sorted({L_snap} | {l for l in ALLOWED_LENGTHS if l > L_snap}):
                cand = _score_candidate((z, x_c, az_c, altitude, d_snap, L_c))
                if cand is None:
                    continue
                if cand['in_bone'] >= LATERAL_REFINE_MIN_IN_BONE and cand['loss'] < best['loss']:
                    best = cand

    z_r, x_r, az_r, alt_r, d_r, L_r = best['pos']
    adopted = not (abs(z_r - z) < 1e-6 and abs(x_r - x) < 1e-6
                   and abs(az_r - azimuth) < 1e-6 and abs(L_r - L_snap) < 1e-6)
    return z_r, x_r, az_r, alt_r, d_r, L_r, adopted, best


def _validate_bounds(lb, ub, name):
    for i, (lo, hi) in enumerate(zip(lb, ub)):
        if lo >= hi:
            raise ValueError(f'PSO bounds invalid for {name}: dim {i} lower {lo} >= upper {hi}')


def get_first_nonzero_y(arr, endplate_plane=None, vbody_mask=None):
    """每個 (z, x) 柱的第一個 bone voxel 的 y（voxel index）。

    endplate_plane 為 {'plane': (a, b, c, d), ...}（法線朝上，c > 0）時，
    入口點 (x, y, z) 落在終板面之上的柱也設 OUTSIDE_VALUE，
    避免從終板上方的柱選入射點。

    vbody_mask（(z, y, x) 0/1，VBODY 椎體 mask）提供時，
    入口點落在椎體上的柱設 OUTSIDE_VALUE，避免螺絲入口點放在椎體上：
    入口 voxel 本身在 VBODY 內、或往前（+y，椎體方向） VBODY_ENTRY_EDGE
    個 voxel 內進入 VBODY（椎體後側皮質邊緣，mask 外 0.5~2mm 的交界帶）
    都算。真正的後側要素（椎弓根/椎板/關節突）厚度遠大於該邊緣，
    且與椎體之間有椎間孔間隔，不會被誤排。
    """
    OUTSIDE_VALUE = -100

    # 1. Create a boolean mask where elements are non-zero
    mask = arr != 0
    
    # 2. Find the index of the first True value along the Y axis (axis=1)
    y_indices = np.argmax(mask, axis=1)
    
    # 3. Edge Case Handling: If a whole (x, z) column is zero, argmax returns 0.
    # We need to distinguish this from an actual non-zero value at index 0.
    has_nonzero = np.any(mask, axis=1)
    
    # 4. Replace indices where there were no non-zeros with a sentinel value (e.g., -1)
    y_indices = np.where(has_nonzero, y_indices, OUTSIDE_VALUE)
    
    y_indices = np.where(y_indices < arr.shape[1] * .1, OUTSIDE_VALUE, y_indices)
    y_indices = np.where(y_indices > arr.shape[1] * .4, OUTSIDE_VALUE, y_indices)

    # 5. 入口點在終板面上方（a*x + b*y + c*z - d > 0）的柱設 OUTSIDE_VALUE
    if endplate_plane is not None and y_indices.max() >= 0:
        a, b, c, d = endplate_plane['plane']
        z_grid, x_grid = np.meshgrid(np.arange(arr.shape[0]), np.arange(arr.shape[2]),
                                     indexing='ij')
        above = (y_indices >= 0) & (a * x_grid + b * y_indices + c * z_grid > d)
        y_indices = np.where(above, OUTSIDE_VALUE, y_indices)

    # 6. 入口點落在椎體上的柱設 OUTSIDE_VALUE（入口不放椎體）：
    #    入口 voxel 本身在 VBODY 內，或往前 VBODY_ENTRY_EDGE 個 voxel 內進入
    #    VBODY（椎體後側皮質邊緣，mask 外 0.5~2mm 的分割界線帶）都排除。
    if vbody_mask is not None:
        vb = np.asarray(vbody_mask) > 0
        valid = y_indices >= 0
        y_i = np.where(valid, y_indices, 0).astype(np.int64)
        z_grid, x_grid = np.meshgrid(np.arange(arr.shape[0]), np.arange(arr.shape[2]),
                                     indexing='ij')
        on_vbody = np.zeros(y_indices.shape, dtype=bool)
        for k in range(VBODY_ENTRY_EDGE + 1):
            y_k = np.minimum(y_i + k, arr.shape[1] - 1)
            on_vbody |= valid & vb[z_grid, y_k, x_grid]
        n_vb = int(on_vbody.sum())
        if n_vb:
            y_indices = np.where(on_vbody, OUTSIDE_VALUE, y_indices)
            print(f"[Y-INDEX] VBODY entry excluded: {n_vb} columns whose entry "
                  f"is on/within {VBODY_ENTRY_EDGE} vox of VBODY removed from "
                  f"entry surface")

    return y_indices.astype(np.float32)

def constraint_y(x, y_indices):
    return y_indices[round(x[0]), round(x[1])]

def run_pso_torch_xfr(
    label_str: str,
    image1_path: str,
    image2_path: str,
    image3_path: str,
    folder: str,
    swarm_size: int,
    max_iter: int,
    spacing: list,
    CBT: bool,
    device: torch.device,
    optimize_size: bool = True,
    grid=None,
    debug=False,
    omega = 0.9,
    side: str = 'both',
    level: str = None,
    patient_id: str = None,
    side_dir: str = None,
    run_id: str = None,
):
    """
    Main function to run PSO. 
    如果 optimize_size=True，diameter 和 length 也會被最佳化
    如果 optimize_size=False，使用預設值（向後兼容）
    """
    start_time = time.time()

    # Use global references
    global image1_array, image2_array, image2_shape, image3_array
    global diameter, length  # 這些現在只用於非最佳化模式
    global spine_tensor, cortical_tensor, spine_roi_tensor

    # Load images
    image1 = sitk.ReadImage(image1_path)
    image2 = sitk.ReadImage(image2_path)
    image3 = sitk.ReadImage(image3_path)
    image1_array = sitk.GetArrayFromImage(image1)
    image2_array = sitk.GetArrayFromImage(image2)
    image3_array = sitk.GetArrayFromImage(image3)
    image2_shape = image2_array.shape
    image_shape = image2_shape

    # Move arrays to torch
    cortical_tensor = torch.from_numpy(image1_array).to(device=device, dtype=torch.uint8)
    spine_tensor = torch.from_numpy(image2_array).to(device=device, dtype=torch.uint8)
    spine_roi_tensor = torch.from_numpy(image3_array).to(device=device, dtype=torch.uint8)

    # 建立明確的優化上下文（取代舊的跨檔案 global 注入；狀態全部顯式傳入目標函數）
    ctx = OptimizationContext(
        cortical_tensor=cortical_tensor,
        spine_tensor=spine_tensor,
        spine_roi_tensor=spine_roi_tensor,
        image1_array=image1_array,
        image2_array=image2_array,
        image3_array=image3_array,
        image2_shape=image2_shape,
        spacing=spacing,
        device=device,
        grid=grid,
        diameter=diameter if not optimize_size else None,
        length=length if not optimize_size else None,
    )

    if not CBT:
        azi = azimuth_rotation(image2_path)
        res = analyze_vertebral_tilt_contour(image2_path, edge_type='superior', show_plot=False, debug=False)
        alt = res['superior']['tilt_angle_deg']

    # ===== 平面：影像已由 xfr_preprocess 對齊旋轉到椎體基準系 =====
    # （鏡稱面法線 -> +x、上終板 normal y=0 / z>0），az/alt 範圍改用固定
    # 約束（見下方 CBT bounds），不再逐椎以 theta_v / tau 重新錨定。
    # 此處平面仅供：
    #   - 棘突移除（segment_spinous_process，sym）
    #   - 入口面終板上方剪除（get_first_nonzero_y，endplate）
    #   - VBODY 椎體分割（segment_vertebral_body，sym + endplate）
    #   - 對齊健全性檢查：基準系下 theta_v ≈ 0、tau_y ≈ 0（偏大 = 預處理失效）
    sym_plane = best_symmetry_plane(image2_array)

    # 入口面 y_indices 取「每個 (z,x) 柱第一個 bone voxel（最後側）」，
    # 中線柱會落在棘突上（入口太靠內後）。先把棘突（鏡稱面中線後側，
    # 見 segment_spinous_process）從 image2_array 移除再取 surface，
    # 讓 y_indices 永不落在棘突上（loss 用 spine_tensor，不受影響）。
    sp_mask, sp_th, sp_info = segment_spinous_process(image2_array, sym_plane)

    # 上終板面與椎體都在「完整 mask（SP 移除前）」上計算：
    #   - 終板面由前側頂面擬合，SP 移除不改變平面；
    #   - 椎體與 render_bone_figure 的 gold 顯示完全同 input / 同參數，
    #     確保顯示出來的椎體就是 loss 裡 VBODY 獎勵的區域。
    # 棘突缺如（laminectomy，mode='no_spinous'）時不該把殘留後側要素
    # 當「棘突」移除（會鏟進椎體後側），入口面維持完整 mask。
    endplate_plane = best_upper_endplate_plane(image2_array)

    # VBODY 評分獎勵：兩平面（鏡稱面 + 上終板）切出的椎體 mask
    vb_mask_np, vb_th, vb_info = segment_vertebral_body(image2_array, sym_plane,
                                                        endplate_plane, sp_th, sp_info['mode'])
    vbody_tensor = None
    if vb_mask_np is not None and vb_mask_np.any():
        vbody_tensor = torch.from_numpy(vb_mask_np.astype(np.uint8)).to(device=device)
        print(f"[VBODY-SCORE] n={vb_info['n_vb']} "
              f"({100.0 * vb_info['n_vb'] / max(int(image2_array.sum()), 1):.1f}% of bone) "
              f"AP<{vb_info['ap_thresh']:.1f} mode={vb_info['mode']} -> added to loss")
        if vb_info['mode'] == 'quantile':
            print(f"[VBODY-SCORE] WARNING: 未找到體/弓後側谷底，閾值退回 55 百分位 "
                  f"（可能切進椎體內），建議人工核對該 level 的椎體邊界")
    else:
        print(f"[VBODY-SCORE] skipped: {vb_info['mode']}")
    ctx.vbody_tensor = vbody_tensor

    if sp_info['mode'] == 'no_spinous':
        top_off = f"{sp_info['top_off']:.1f}" if sp_info.get('top_off') is not None else 'n/a'
        print(f"[NO-SP] 中線後側缺如（先前 laminectomy / 棘突切除）: "
              f"deficit={sp_info['deficit']:.1f} vox ({sp_info['deficit'] * 0.5:.1f} mm), "
              f"rear3={sp_info['rear3']} vox, top_off={top_off} vox "
              f"-> 不移除 SP（完整 mask 取入口面），椎體用放寬後側谷底切分")
    elif sp_mask is not None and sp_mask.any():
        n_sp = int(sp_mask.sum())
        image2_array[sp_mask] = 0
        print(f"[Y-INDEX] spinous process removed from entry surface: {n_sp} vox "
              f"(band=+/-{sp_info['band_w']:.1f} voxel, AP>={sp_info['ap_thresh']:.1f}, "
              f"mode={sp_info['mode']})")

    y_indices = get_first_nonzero_y(image2_array, endplate_plane, vb_mask_np)

    objective_fn = make_objective_function_xfr(ctx, y_indices)

    # 對齊健全性檢查（基準系下 theta_v ≈ 0、tau_y ≈ 0；不參與 bounds 計算）
    s_nx, s_ny, s_nz = sym_plane['normal']
    theta_v = float(np.degrees(np.arctan2(s_ny, s_nx)))
    if endplate_plane is not None:
        e_nx, e_ny, e_nz = endplate_plane['normal']
        tau_y = float(np.degrees(np.arctan2(e_ny, e_nz)))
        tau_x = float(np.degrees(np.arctan2(e_nx, e_nz)))
    else:
        tau_y, tau_x = 0.0, 0.0
    print(f"[PLANE] mirror  : {sym_plane['plane'][0]:+.3f}x {sym_plane['plane'][1]:+.3f}y "
          f"{sym_plane['plane'][2]:+.3f}z = {sym_plane['plane'][3]:.1f}"
          f"  (theta_v={theta_v:+.2f} deg, mirror ratio={sym_plane['ratio']:.3f})")
    if endplate_plane is not None:
        p = endplate_plane['plane']
        print(f"[PLANE] endplate: {p[0]:+.3f}x {p[1]:+.3f}y {p[2]:+.3f}z = {p[3]:.1f}"
              f"  (tau_y={tau_y:+.2f} deg, tau_x={tau_x:+.2f} deg, "
              f"inlier={endplate_plane['inlier_ratio']:.2f})")
    else:
        print("[PLANE] endplate: 資料不足，終板入口剪除停用")

    # flat_min_index = np.argmin(y_indices)
    # z_border, x_border = np.unravel_index(flat_min_index, y_indices.shape)

    # x/z 搜索空間邊界：把 VBODY 椎體投影到 xz 平面，取該投影的 bounding box，
    # 再在 x 向切 L/R 兩 band + 中央缺口、z 向留 10%~90%（見下方 CBT bounds）。
    # CBT 入口點應落在椎體上（左/右 band），而非跨整段骨頭：整段骨頭包含
    # 肋、後側要素、橫突等極端，x 範圍比椎體寬，會把 L/R band 往外推。
    # （不取單一行 / 不用鏡稱對稱軸的原因同前，舊註保留於下）。
    # VBODY mask 為 (z,y,x)；np.any(..., axis=1) 折疊 y 得 xz 投影 (z,x)。
    #   x 範圍 = 有 VBODY 的欄（投影 axis=0 是 z，沿 z 做 any）
    #   z 範圍 = 有 VBODY 的行（投影 axis=1 是 x，沿 x 做 any）
    # VBODY 分割失敗（None / 全 0）時退回整段骨頭 x/z 範圍。
    if vb_mask_np is not None and vb_mask_np.any():
        vb_xz = np.any(vb_mask_np, axis=1)          # (z, x) 投影
        x1 = int(np.where(vb_xz.any(axis=0))[0][0])
        x2 = int(np.where(vb_xz.any(axis=0))[0][-1])
        z1 = int(np.where(vb_xz.any(axis=1))[0][0])
        z2 = int(np.where(vb_xz.any(axis=1))[0][-1])
        print(f"[BOUNDS] VBODY xz projection: z[{z1},{z2}] x[{x1},{x2}] "
              f"(z_height={z2 - z1}, x_width={x2 - x1})")
    else:
        print("[BOUNDS] VBODY unavailable, falling back to whole-bone xz range")
        x_with_nonzero = np.where(np.any(image2_array, axis=(0, 1)))[0]
        x1 = int(x_with_nonzero[0])
        x2 = int(x_with_nonzero[-1])
        z1 = int(z_with_nonzero[0])
        z2 = int(z_with_nonzero[-1])
    x_width = x2 - x1
    z_height = z2 - z1

    # print(x1,x2)
    # exit()

    # print(x_border, z_border)
    # exit()
    # import sys
    # import numpy
    # numpy.set_printoptions(threshold=sys.maxsize)
    # print(y_indices)
    # exit()

    # 設定基本的 bounds
    if CBT == True:    
        # z_bounds = (0, image_shape[0]-1)
        # z_bounds = (z1, (z1+z2)/2)
        # z_bounds = (.1*image_shape[0], .8*image_shape[0])
        z_bounds = (0, z1+z_height*.7)

        # x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1)
        # x_bounds_left  = (0, image_shape[2]/2 - image_shape[2]/10 - 1)
        # x_bounds_right = (x2, image_shape[2]*.9)
        # x_bounds_left  = (image_shape[2]*.1, x1)
        x_bounds_left  = (x1+x_width*.1, x1+x_width*.4)
        x_bounds_right = (x1+x_width*.6, x1+x_width*.9)

        # 脊椎若被體積邊界切到（真正偏心、骨頭貼著左/右邊緣），
        # 對應那側的 x band 下限會 >= 上限，PSO 會丟 "upper-bound must be greater"。
        # 出錯時 clamp 成同側最小寬度（5% 寬度）的合法 band。
        min_band = .05 * image_shape[2]
        if x_bounds_left[1] <= x_bounds_left[0]:
            x_bounds_left = (x_bounds_left[0], x_bounds_left[0] + min_band)
        if x_bounds_right[0] >= x_bounds_right[1]:
            x_bounds_right = (x_bounds_right[1] - min_band, x_bounds_right[1])

        # 固定約束（影像已由 xfr_preprocess 對齊旋轉到椎體基準系：鏡稱面法線 +x、
        # 終板 normal y=0 / z>0；基準系下 theta_v≈0、tau_y≈0，舊式逐椎錨定
        # (98+theta_v, 105+theta_v) / (60+tau_y∓tau_x*sin_mid, 70+tau_y∓tau_x*sin_mid)
        # 不再需要）：
        #   azimuth (Lateral)：az=90° 為 AP 直向，每側向外發散 8~20°
        #            L = 90 + (8~20) = 98~110、R = 90 - (8~20) = 70~82
        #   altitude (Cephalad)：相對終板面 25~30°（+z 極角 90 - 25~30）= 60~65
        # 舊 2D 版本（以輪廓角 azi / 矢狀面傾斜 alt 平移固定範圍），保留供對照：
        # azimuth_bounds_l = ((98-azi), (120-azi))
        # azimuth_bounds_r = ((60-azi), (82-azi))
        # altitude_bounds = ((60-alt), (70-alt))
        azimuth_bounds_l = (98, 105)
        azimuth_bounds_r = (75, 82)
        altitude_bounds_l = (60, 65)
        altitude_bounds_r = (60, 65)

    else:      
        z_bounds = (0, image_shape[0] - 1)   
        y_bounds = (image_shape[1]/5, image_shape[1]/2 - 1)   
        x_bounds_left = (0, image_shape[2]/2 - image_shape[2]/10 - 1)   
        x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1)
        azimuth_bounds_l = (60-azi, 90-azi) 
        azimuth_bounds_r = (90-azi, 120-azi) 
        altitude_bounds = (65-alt, 80-alt)
        # 非 CBT 分支維持單一式，兩側同用
        altitude_bounds_l = altitude_bounds
        altitude_bounds_r = altitude_bounds

    def eval_overlap_from_position(pos, side: str, optimize_size: bool,
                               spine_tensor: torch.Tensor,
                               image_shape, spacing):
        """
        根據 PSO 給的 position 生成 cylinder mask，再算 overlap ratio
        side: "L" or "R" 只是方便 debug
        """
        if optimize_size:
            # d, L = snap_to_discrete_values(pos[5], pos[6])
            d, L = snap_to_discrete_values_xfr(pos[5], pos[6])
            params_5 = pos[:5]
        else:
            d, L = diameter, length
            params_5 = pos

        cyl_mask = generate_cylinder_n_torch(
            d, L,
            params_5[0], params_5[1], params_5[2],
            params_5[3], params_5[4],
            image_shape, spacing, device, grid
        )

        overlap = compute_overlap_ratio_from_cylinder_mask(cyl_mask, spine_tensor)
        return overlap, d, L

    if optimize_size:
        # 模式 1：優化 diameter 和 length
        print("=== 最佳化模式：最佳化位置、角度、直徑和長度 ===")
        
        # 設定 diameter 和 length 的 bounds（連續範圍）
        diameter_bounds = (min(ALLOWED_DIAMETERS), max(ALLOWED_DIAMETERS)*1.01)
        length_bounds = (min(ALLOWED_LENGTHS), max(ALLOWED_LENGTHS)*1.01)
        
        # bounds 現在有 6 個參數 [z, x, az, alt, d, L]（y 由 y_indices 取）；
        # az/alt 兩側用同一組固定約束，x 帶依 L/R 分開
        lb_l = [z_bounds[0], x_bounds_left[0], azimuth_bounds_l[0], 
                altitude_bounds_l[0], diameter_bounds[0], length_bounds[0]]
        ub_l = [z_bounds[1], x_bounds_left[1], azimuth_bounds_l[1], 
                altitude_bounds_l[1], diameter_bounds[1], length_bounds[1]]

        lb_r = [z_bounds[0], x_bounds_right[0], azimuth_bounds_r[0], 
                altitude_bounds_r[0], diameter_bounds[0], length_bounds[0]]
        ub_r = [z_bounds[1], x_bounds_right[1], azimuth_bounds_r[1], 
                altitude_bounds_r[1], diameter_bounds[1], length_bounds[1]]
        
    else:
        # 模式 2：固定 diameter 和 length（向後兼容）
        print("=== 固定尺寸模式：最佳化位置和角度 ===")
        # 使用預設值（需要在調用時提供）
        diameter = 4.5  # 或從參數傳入
        length = 45     # 或從參數傳入
        
        lb_l = [z_bounds[0], y_bounds[0], x_bounds_left[0], azimuth_bounds_l[0], altitude_bounds_l[0]]
        ub_l = [z_bounds[1], y_bounds[1], x_bounds_left[1], azimuth_bounds_l[1], altitude_bounds_l[1]]

        lb_r = [z_bounds[0], y_bounds[0], x_bounds_right[0], azimuth_bounds_r[0], altitude_bounds_r[0]]
        ub_r = [z_bounds[1], y_bounds[1], x_bounds_right[1], azimuth_bounds_r[1], altitude_bounds_r[1]]

    if True or debug:
        for b in (lb_l, ub_l, lb_r, ub_r):
            print('[' + ', '.join(f'{v:10.2f}' for v in b) + ']')

    best_loss_l = float('inf')
    best_loss_r = float('inf')
    best_position_l = None
    best_position_r = None

    # L/R 是兩次獨立 PSO（bounds 不同、目標函數共用），互不依賴：
    # side='both' 維持原行為（同一次呼叫先 L 後 R）；
    # side='L'/'R' 只跑該側，讓兩側可排到不同 GPU worker。
    # 單側模式的合併輸出（3D 圖 + CSV）由「較晚完成」的一側在
    # 下方「單側收尾」段觸發。
    def _side_bounds(s_i):
        if s_i == 'L':
            return lb_l, ub_l, azimuth_bounds_l, x_bounds_left
        return lb_r, ub_r, azimuth_bounds_r, x_bounds_right

    def _run_one_side(s_i):
        lb_s, ub_s, az_bounds_s, x_bounds_s = _side_bounds(s_i)
        tag_cn = '左側' if s_i == 'L' else '右側'
        tag = 'LEFT' if s_i == 'L' else 'RIGHT'
        print(f"\n=== {label_str} {tag_cn} ===")
        _validate_bounds(lb_s, ub_s, f'{label_str} {s_i}')
        position, loss = pso(objective_fn, lb_s, ub_s,
            # ieqcons=[constraint_y],
            swarmsize=swarm_size,
            omega=omega,
            maxiter=max_iter, debug=debug)

        # 如果需要 retry（loss > 0）：重跑 PSO 取較好者
        # （原 L/R 各一份註解版 retry，此處合併為一式）
        # max_retries = 0
        # retries = 0
        # while loss > 0 and retries < max_retries:
        #     position, loss = pso(objective_fn, lb_s, ub_s,
        #                          swarmsize=swarm_size, maxiter=max_iter)
        #     retries += 1

        z, x, azimuth, altitude, diameter, length = position
        az_pso, x_pso, L_pso = azimuth, x, length
        y = y_indices[round(z), round(x)]

        z, x, azimuth, altitude, diameter, length, adopted, ref = refine_lateral_longer(
            z, x, azimuth, altitude, diameter, length,
            s_i, az_bounds_s, x_bounds_s, y_indices,
            image_shape, spacing, device, grid,
            cortical_tensor, spine_tensor,
            vbody_tensor,
        )
        y = y_indices[round(z), round(x)]
        if adopted:
            loss = ref['loss']
            print(f"[{tag}] lateral-refine: az {az_pso:.2f} -> {azimuth:.2f}, x {x_pso:.2f} -> {x:.2f}, "
                  f"L {L_pso:.2f} -> {length:.2f}, in-bone {ref['in_bone']*100:.1f}%")
        else:
            print(f"[{tag}] lateral-refine: no improvement")

        position = z, y, x, azimuth, altitude, diameter, length

        overlap_s, d_snap, L_snap = eval_overlap_from_position(
            position, s_i, optimize_size, spine_tensor, image_shape, spacing
        )
        print(f"[{tag}] overlap: {overlap_s*100:.1f}%")

        if optimize_size:
            print(f"[{tag}] Position: {position[:5]}")
            print(f"[{tag}] Diameter: {d_snap} mm (raw: {position[5]:.2f})")
            print(f"[{tag}] Length: {L_snap} mm (raw: {position[6]:.2f})")
            print(f"[{tag}] Loss: {loss}\n")
            best_pos = list(position[:5]) + [d_snap, L_snap]
        else:
            print(f"[{tag}] Position: {position}")
            best_pos = position
        return best_pos, loss, overlap_s

    sides = ('L', 'R') if side == 'both' else (side,)
    side_time = {}
    for s_i in sides:
        t0 = time.time()
        pos_s, loss_s, _ = _run_one_side(s_i)
        side_time[s_i] = time.time() - t0
        if s_i == 'L':
            best_position_l, best_loss_l = pos_s, loss_s
        else:
            best_position_r, best_loss_r = pos_s, loss_s

    end_time = time.time()
    total_time = end_time - start_time

    # 提取最終的 diameter 和 length
    if optimize_size:
        final_diameter_l = best_position_l[5] if best_position_l is not None else float('nan')
        final_length_l = best_position_l[6] if best_position_l is not None else float('nan')
        final_diameter_r = best_position_r[5] if best_position_r is not None else float('nan')
        final_length_r = best_position_r[6] if best_position_r is not None else float('nan')
    else:
        final_diameter_l = diameter
        final_length_l = length
        final_diameter_r = diameter
        final_length_r = length

    def _plot_combined(d_l, l_l, d_r, l_r, pos_l, pos_r, total_t):
        # volume / level 由 image2_path 反推（…/{vol}/rotated/{level}_*.nii.gz）
        render_bone_figure(
            None, None,
            spine_tensor,
            cortical_tensor,
            folder,
            spacing=spacing,
            way='CBT' if CBT else 'TPS',
            best_position_l=pos_l, best_position_r=pos_r,
            diameter_l=d_l, length_l=l_l,
            diameter_r=d_r, length_r=l_r,
            image2_path=image2_path,
            device=device, grid=grid,
            swarm_size=swarm_size, max_iter=max_iter, total_time=total_t,
        )

    if side == 'both':
        print(f"\n=== {label_str} 最終結果 ===")
        print(f"Left - Diameter: {final_diameter_l} mm, Length: {final_length_l} mm")
        print(f"Right - Diameter: {final_diameter_r} mm, Length: {final_length_r} mm")
        _plot_combined(final_diameter_l, final_length_l,
                       final_diameter_r, final_length_r,
                       best_position_l, best_position_r, total_time)
        return best_position_l, best_loss_l, best_position_r, best_loss_r, total_time

    # ---- 單側收尾：寫該側結果檔；兩側都完成時由較晚的一側觸發合併輸出 ----
    # 兩側可能在不同 GPU worker：各自把結果寫 <level>_<side>.json
    # （先寫 tmp 再 os.replace，對端讀到的一定是完整檔）。先完成者看不到
    # 對端檔就跳過；後完成者看到兩側齊了、搶到 plot lock（O_EXCL，
    # 確保合併輸出只跑一次）才載入對端結果跑 render_bone_figure
    # （3D 圖 + CSV 兩行，與 'both' 模式相同）。json / lock 保留供事後
    # 檢查；若該側流程死在 plotting 中段，該 (volume, level) 重跑即可
    # （run_id 是新的一次，不會互相干擾）。
    other = 'R' if side == 'L' else 'L'
    own_pos = best_position_l if side == 'L' else best_position_r
    own_d = final_diameter_l if side == 'L' else final_diameter_r
    own_L = final_length_l if side == 'L' else final_length_r
    own_cn = 'Left' if side == 'L' else 'Right'
    side_cn = '左' if side == 'L' else '右'
    print(f"\n=== {label_str} 最終結果（{side_cn}側） ===")
    print(f"{own_cn} - Diameter: {own_d} mm, Length: {own_L} mm")

    missing_pair = [n for n, v in (('level', level), ('patient_id', patient_id),
                                   ('side_dir', side_dir), ('run_id', run_id)) if not v]
    if missing_pair:
        print(f"[SIDE] 合併輸出跳過（未提供 {', '.join(missing_pair)}）")
    else:
        patient_dir = os.path.join(side_dir, run_id, patient_id)
        own_path = os.path.join(patient_dir, f'{level}_{side}.json')
        other_path = os.path.join(patient_dir, f'{level}_{other}.json')
        _makedirs_retry(patient_dir)
        own = {'position': [float(v) for v in own_pos],
               'diameter': float(own_d),
               'length': float(own_L),
               'time': float(side_time[side])}
        tmp_path = f'{own_path}.{os.getpid()}.tmp'
        with open(tmp_path, 'w') as f:
            json.dump(own, f)
        os.replace(tmp_path, own_path)

        if not os.path.isfile(other_path):
            print(f"[SIDE] {other} 側尚未完成，合併輸出（圖 + CSV）待該側完成時觸發")
        else:
            lock_path = os.path.join(patient_dir, f'{level}.plot.lock')
            try:
                fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
                os.close(fd)
            except FileExistsError:
                print(f'[SIDE] 合併輸出已由 {other} 側觸發，跳過')
            else:
                with open(other_path) as f:
                    other_res = json.load(f)
                pos_other = [float(v) for v in other_res['position']]
                if side == 'L':
                    d_l, l_l, pos_l = own_d, own_L, list(own_pos)
                    d_r, l_r, pos_r = other_res['diameter'], other_res['length'], pos_other
                else:
                    d_r, l_r, pos_r = own_d, own_L, list(own_pos)
                    d_l, l_l, pos_l = other_res['diameter'], other_res['length'], pos_other
                print(f"[SIDE] {side} + {other} 兩側完成 -> 合併輸出（圖 + CSV）")
                # total_time 用兩側各自耗時相加（各含一次影像載入/平面計算，
                # 比原同流程 wall time 略大，僅影響 CSV 的時間欄）
                _plot_combined(d_l, l_l, d_r, l_r, pos_l, pos_r,
                               side_time[side] + other_res['time'])

    return best_position_l, best_loss_l, best_position_r, best_loss_r, total_time


def run_pso_torch(
    label_str: str,
    image1_path: str,
    image2_path: str,
    image3_path: str,
    folder: str,
    swarm_size: int,
    max_iter: int,
    spacing: list,
    CBT: bool,
    device: torch.device,
    optimize_size: bool = True,
    grid=None,
    debug=True,
):
    """
    Main function to run PSO. 
    如果 optimize_size=True，diameter 和 length 也會被最佳化
    如果 optimize_size=False，使用預設值（向後兼容）
    """
    start_time = time.time()

    # Use global references
    global image1_array, image2_array, image2_shape, image3_array
    global diameter, length  # 這些現在只用於非最佳化模式
    global spine_tensor, cortical_tensor, spine_roi_tensor

    # Load images
    image1 = sitk.ReadImage(image1_path)
    image2 = sitk.ReadImage(image2_path)
    image3 = sitk.ReadImage(image3_path)
    image1_array = sitk.GetArrayFromImage(image1)
    image2_array = sitk.GetArrayFromImage(image2)
    image3_array = sitk.GetArrayFromImage(image3)
    image2_shape = image2_array.shape
    image_shape = image2_shape

    # Move arrays to torch
    cortical_tensor = torch.from_numpy(image1_array).to(device=device, dtype=torch.uint8)
    spine_tensor = torch.from_numpy(image2_array).to(device=device, dtype=torch.uint8)
    spine_roi_tensor = torch.from_numpy(image3_array).to(device=device, dtype=torch.uint8)

    # 建立明確的優化上下文（取代舊的跨檔案 global 注入；狀態全部顯式傳入目標函數）
    ctx = OptimizationContext(
        cortical_tensor=cortical_tensor,
        spine_tensor=spine_tensor,
        spine_roi_tensor=spine_roi_tensor,
        image1_array=image1_array,
        image2_array=image2_array,
        image3_array=image3_array,
        image2_shape=image2_shape,
        spacing=spacing,
        device=device,
        grid=grid,
        diameter=diameter if not optimize_size else None,
        length=length if not optimize_size else None,
    )
    objective_fn = make_objective_function(ctx)

    azi = azimuth_rotation(image2_path)
    res = analyze_vertebral_tilt_contour(image2_path, edge_type='superior', show_plot=False, debug=False)
    alt = res['superior']['tilt_angle_deg']

    # 設定基本的 bounds
    if CBT == True:    
        z_bounds = (0, image_shape[0] - 1)
        y_bounds = (image_shape[1]/5, image_shape[1]/2 - 1)
        x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1)
        x_bounds_left  = (0, image_shape[2]/2 - image_shape[2]/10 - 1)
        # CBT 參數依據 references/：Santoni 2009 (Spine J 9:366) 冠狀面 25-30° cranial (caudo-cephalad)、
        # 軸狀面自正中線向外 (medial→lateral) ≤30°；Delgado-Fernandez 2017 (ASJ 11:817)、Kim 2022 (SSRR 6:1)
        azimuth_bounds_l = ((98-azi), (120-azi))
        azimuth_bounds_r = ((60-azi), (82-azi))
        altitude_bounds = ((60-alt), (70-alt))

        # xfr
        # z_bounds = (0, image_shape[0] - 1)
        y_bounds = (0, image_shape[1]/2)
        # x_bounds_right = (image_shape[2]/2, image_shape[2] - 1)
        # x_bounds_left  = (0, image_shape[2]/2)
        # azimuth_bounds_l = (90, 135)
        # azimuth_bounds_r = (45, 90)
        # altitude_bounds = (0, 90)
    else:      
        z_bounds = (0, image_shape[0] - 1)   
        y_bounds = (image_shape[1]/5, image_shape[1]/2 - 1)   
        x_bounds_left = (0, image_shape[2]/2 - image_shape[2]/10 - 1)   
        x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1)
        azimuth_bounds_l = (60-azi, 90-azi) 
        azimuth_bounds_r = (90-azi, 120-azi) 
        altitude_bounds = (65-alt, 80-alt)

    def eval_overlap_from_position(pos, side: str, optimize_size: bool,
                               spine_tensor: torch.Tensor,
                               image_shape, spacing):
        """
        根據 PSO 給的 position 生成 cylinder mask，再算 overlap ratio
        side: "L" or "R" 只是方便 debug
        """
        if optimize_size:
            # d, L = snap_to_discrete_values(pos[5], pos[6])
            d, L = snap_to_discrete_values_xfr(pos[5], pos[6])
            params_5 = pos[:5]
        else:
            d, L = diameter, length
            params_5 = pos

        cyl_mask = generate_cylinder_n_torch(
            d, L,
            params_5[0], params_5[1], params_5[2],
            params_5[3], params_5[4],
            image_shape, spacing, device, grid
        )

        overlap = compute_overlap_ratio_from_cylinder_mask(cyl_mask, spine_tensor)
        return overlap, d, L

    if optimize_size:
        # 模式 1：優化 diameter 和 length
        print("=== 最佳化模式：最佳化位置、角度、直徑和長度 ===")
        
        # 設定 diameter 和 length 的 bounds（連續範圍）
        diameter_bounds = (min(ALLOWED_DIAMETERS), max(ALLOWED_DIAMETERS))
        length_bounds = (min(ALLOWED_LENGTHS), max(ALLOWED_LENGTHS))
        
        # bounds 現在有 7 個參數
        lb_l = [z_bounds[0], y_bounds[0], x_bounds_left[0], azimuth_bounds_l[0], 
                altitude_bounds[0], diameter_bounds[0], length_bounds[0]]
        ub_l = [z_bounds[1], y_bounds[1], x_bounds_left[1], azimuth_bounds_l[1], 
                altitude_bounds[1], diameter_bounds[1], length_bounds[1]]

        lb_r = [z_bounds[0], y_bounds[0], x_bounds_right[0], azimuth_bounds_r[0], 
                altitude_bounds[0], diameter_bounds[0], length_bounds[0]]
        ub_r = [z_bounds[1], y_bounds[1], x_bounds_right[1], azimuth_bounds_r[1], 
                altitude_bounds[1], diameter_bounds[1], length_bounds[1]]
        
    else:
        # 模式 2：固定 diameter 和 length（向後兼容）
        print("=== 固定尺寸模式：最佳化位置和角度 ===")
        # 使用預設值（需要在調用時提供）
        diameter = 4.5  # 或從參數傳入
        length = 45     # 或從參數傳入
        
        lb_l = [z_bounds[0], y_bounds[0], x_bounds_left[0], azimuth_bounds_l[0], altitude_bounds[0]]
        ub_l = [z_bounds[1], y_bounds[1], x_bounds_left[1], azimuth_bounds_l[1], altitude_bounds[1]]

        lb_r = [z_bounds[0], y_bounds[0], x_bounds_right[0], azimuth_bounds_r[0], altitude_bounds[0]]
        ub_r = [z_bounds[1], y_bounds[1], x_bounds_right[1], azimuth_bounds_r[1], altitude_bounds[1]]

    if debug:
        print(lb_l)
        print(ub_l)
        print(lb_r)
        print(ub_r)

    best_loss_l = float('inf')
    best_loss_r = float('inf')
    best_position_l = None
    best_position_r = None

    # Left side optimization
    print("\n=== 左側 ===")
    position_l, loss_l = pso(objective_fn, lb_l, ub_l, swarmsize=swarm_size, maxiter=max_iter, debug=debug)

    overlap_l, diameter_l, length_l = eval_overlap_from_position(
        position_l, "L", optimize_size, spine_tensor, image_shape, spacing
    )
    print(f"[LEFT] overlap: {overlap_l*100:.1f}%")

    if optimize_size:
        print(f"[LEFT] Position: {position_l[:5]}")
        print(f"[LEFT] Diameter: {diameter_l} mm (raw: {position_l[5]:.2f})")
        print(f"[LEFT] Length: {length_l} mm (raw: {position_l[6]:.2f})")
        best_position_l = list(position_l[:5]) + [diameter_l, length_l]
    else:
        print(f"[LEFT] Position: {position_l}")
        best_position_l = position_l

    best_loss_l = loss_l
    best_overlap_l = overlap_l   # 新增
    
    # max_retries = 0
    # retries = 0

    # 左側 retry：loss 要 <=0 且 overlap >= 0.5 才算過關
    # while (best_loss_l > 0 or best_overlap_l < OVERLAP_THRESH) and retries < max_retries:
        # position_l, loss_l = pso(objective_fn, lb_l, ub_l, swarmsize=swarm_size, maxiter=max_iter)
        # overlap_l, diameter_l, length_l = eval_overlap_from_position(
            # position_l, "L", optimize_size, spine_tensor, image_shape, spacing
        # )

    # 只要找到更好的 loss（或你想用 loss+overlap 綜合排序也行）就更新 best
    # 安全版本：優先選「合格解」；沒有合格解時才用 loss 最小的當備案
        # candidate_pos = (list(position_l[:5]) + [diameter_l, length_l]) if optimize_size else position_l

        # candidate_ok = is_solution_ok(loss_l, overlap_l, OVERLAP_THRESH)
        # best_ok = is_solution_ok(best_loss_l, best_overlap_l, OVERLAP_THRESH)

        # if candidate_ok and (not best_ok or loss_l < best_loss_l):
            # best_position_l = candidate_pos
            # best_loss_l = loss_l
            # best_overlap_l = overlap_l
            # print(f"[LEFT][retry {retries+1}] ✅ ok | loss={loss_l:.4f}, overlap={overlap_l*100:.1f}%")
        # elif (not best_ok) and (loss_l < best_loss_l):
            # best 還不合格時，先用更小 loss 的當暫存（至少越來越好）
            # best_position_l = candidate_pos
            # best_loss_l = loss_l
            # best_overlap_l = overlap_l
            # print(f"[LEFT][retry {retries+1}] ⚠️ not ok | loss improved={loss_l:.4f}, overlap={overlap_l*100:.1f}%")
        # else:
            # print(f"[LEFT][retry {retries+1}] ❌ no improve | loss={loss_l:.4f}, overlap={overlap_l*100:.1f}%")

        # retries += 1
    
    # Right side optimization
    print("\n=== 右側 ===")
    position_r, loss_r = pso(objective_fn, lb_r, ub_r, swarmsize=swarm_size, maxiter=max_iter, debug=debug)
    overlap_r, diameter_r, length_r = eval_overlap_from_position(
        position_r, "R", optimize_size, spine_tensor, image_shape, spacing
    )
    print(f"[RIGHT] overlap: {overlap_r*100:.1f}%")

    if optimize_size:
        # diameter_r, length_r = snap_to_discrete_values(position_r[5], position_r[6])
        diameter_r, length_r = snap_to_discrete_values_xfr(position_r[5], position_r[6])
        print(f"[RIGHT] Position: {position_r[:5]}")
        print(f"[RIGHT] Diameter: {diameter_r} mm (raw: {position_r[5]:.2f})")
        print(f"[RIGHT] Length: {length_r} mm (raw: {position_r[6]:.2f})")
        print(f"[RIGHT] Loss: {loss_r}\n")
        
        best_position_r = list(position_r[:5]) + [diameter_r, length_r]
    else:
        print(f"[RIGHT] Position: {position_r}")
        print(f"[RIGHT] Loss: {loss_r}\n")
        best_position_r = position_r
    
    best_loss_r = loss_r
    best_overlap_r = overlap_r 
    
    # 如果需要 retry（loss > 0）
    # max_retries = 10
    # retries = 0

    # while (best_loss_r > 0 or best_overlap_r < OVERLAP_THRESH) and retries < max_retries:
        # position_r, loss_r = pso(objective_fn, lb_r, ub_r, swarmsize=swarm_size, maxiter=max_iter)
        # overlap_r, diameter_r, length_r = eval_overlap_from_position(
            # position_r, "R", optimize_size, spine_tensor, image_shape, spacing
        # )

        # 只要找到更好的 loss（或你想用 loss+overlap 綜合排序也行）就更新 best
        # 這裡給你一個更安全的版本：優先選「合格解」；沒有合格解時才用 loss 最小的當備案
        # candidate_pos = (list(position_r[:5]) + [diameter_r, length_r]) if optimize_size else position_r

        # candidate_ok = is_solution_ok(loss_r, overlap_r, OVERLAP_THRESH)
        # best_ok = is_solution_ok(best_loss_r, best_overlap_r, OVERLAP_THRESH)

        # if candidate_ok and (not best_ok or loss_r < best_loss_r):
            # best_position_r = candidate_pos
            # best_loss_r = loss_r
            # best_overlap_r = overlap_r
            # print(f"[RIGHT][retry {retries+1}] ✅ ok | loss={loss_r:.4f}, overlap={overlap_r*100:.1f}%")
        # elif (not best_ok) and (loss_r < best_loss_r):
            # best 還不合格時，先用更小 loss 的當暫存（至少越來越好）
            # best_position_r = candidate_pos
            # best_loss_r = loss_r
            # best_overlap_r = overlap_r
            # print(f"[RIGHT][retry {retries+1}] ⚠️ not ok | loss improved={loss_r:.4f}, overlap={overlap_r*100:.1f}%")
        # else:
            # print(f"[RIGHT][retry {retries+1}] ❌ no improve | loss={loss_r:.4f}, overlap={overlap_r*100:.1f}%")

        # retries += 1
    
    end_time = time.time()
    total_time = end_time - start_time

    # 提取最終的 diameter 和 length
    if optimize_size:
        final_diameter_l = best_position_l[5]
        final_length_l = best_position_l[6]
        final_diameter_r = best_position_r[5]
        final_length_r = best_position_r[6]
        
        print(f"\n=== 最終結果 ===")
        print(f"Left - Diameter: {final_diameter_l} mm, Length: {final_length_l} mm")
        print(f"Right - Diameter: {final_diameter_r} mm, Length: {final_length_r} mm")
    else:
        final_diameter_l = diameter
        final_length_l = length
        final_diameter_r = diameter
        final_length_r = length

    render_bone_figure(
        None, None,
        spine_tensor,
        cortical_tensor,
        folder,
        spacing=spacing,
        way='CBT' if CBT else 'TPS',
        best_position_l=best_position_l, best_position_r=best_position_r,
        diameter_l=final_diameter_l, length_l=final_length_l,
        diameter_r=final_diameter_r, length_r=final_length_r,
        image2_path=image2_path,
        device=device, grid=grid,
        swarm_size=swarm_size, max_iter=max_iter, total_time=total_time,
    )
    
    return best_position_l, best_loss_l, best_position_r, best_loss_r, total_time

import time
import numpy as np
import SimpleITK as sitk
import torch
from scipy.optimize import differential_evolution
from scipy.optimize import minimize
from imaging.orientation import azimuth_rotation, analyze_vertebral_tilt_contour
from config.constant import ALLOWED_DIAMETERS, ALLOWED_LENGTHS
from core.cylinder import generate_cylinder_n_torch, snap_to_discrete_values, create_coordinate_grid
from core.scoring import compute_overlap_ratio_from_cylinder_mask, is_solution_ok
from config.constant import OVERLAP_THRESH
from visualization.res_bone_figure import render_bone_figure

def run_de_torch(
    label_str: str,
    image1_path: str,
    image2_path: str,
    image3_path: str,
    folder: str,
    swarm_size: int,
    max_iter: int,
    spacing: list,
    CBT: bool,
    device: torch.device,
    optimize_size: bool = True,
    grid=None
):
    """
    使用 Differential Evolution (DE) 進行最佳化
    """
    start_time = time.time()

    global image1_array, image2_array, image2_shape, image3_array
    global diameter, length  
    global spine_tensor, cortical_tensor, spine_roi_tensor

    image1 = sitk.ReadImage(image1_path)
    image2 = sitk.ReadImage(image2_path)
    image3 = sitk.ReadImage(image3_path)
    image1_array = sitk.GetArrayFromImage(image1)
    image2_array = sitk.GetArrayFromImage(image2)
    image3_array = sitk.GetArrayFromImage(image3)
    image2_shape = image2_array.shape
    image_shape = image2_shape

    cortical_tensor = torch.from_numpy(image1_array).to(device=device, dtype=torch.uint8)
    spine_tensor = torch.from_numpy(image2_array).to(device=device, dtype=torch.uint8)
    spine_roi_tensor = torch.from_numpy(image3_array).to(device=device, dtype=torch.uint8)

    # 建立明確的優化上下文（取代舊的跨檔案 global 注入；狀態全部顯式傳入目標函數）
    ctx = OptimizationContext(
        cortical_tensor=cortical_tensor,
        spine_tensor=spine_tensor,
        spine_roi_tensor=spine_roi_tensor,
        image1_array=image1_array,
        image2_array=image2_array,
        image3_array=image3_array,
        image2_shape=image2_shape,
        spacing=spacing,
        device=device,
        grid=grid,
        diameter=diameter if not optimize_size else None,
        length=length if not optimize_size else None,
    )
    objective_fn = make_objective_function(ctx)

    azi = azimuth_rotation(image2_path)
    res = analyze_vertebral_tilt_contour(image2_path, edge_type='superior', show_plot=False, debug=False)
    alt = res['superior']['tilt_angle_deg']

    if CBT == True:    
        z_bounds = (0, image_shape[0] - 1)
        y_bounds = (image_shape[1]/5, image_shape[1]/2 - 1)
        x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1)
        x_bounds_left  = (0, image_shape[2]/2 - image_shape[2]/10 - 1)
        # CBT 參數依據 references/：Santoni 2009 (Spine J 9:366) 冠狀面 25-30° cranial (caudo-cephalad)、
        # 軸狀面自正中線向外 (medial→lateral) ≤30°；Delgado-Fernandez 2017 (ASJ 11:817)、Kim 2022 (SSRR 6:1)
        azimuth_bounds_l = ((98-azi), (120-azi))
        azimuth_bounds_r = ((60-azi), (82-azi))
        altitude_bounds = ((60-alt), (70-alt))
    else:      
        z_bounds = (0, image_shape[0] - 1)   
        y_bounds = (image_shape[1]/5, image_shape[1]/2 - 1)   
        x_bounds_left = (0, image_shape[2]/2 - image_shape[2]/10 - 1)   
        x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1)
        azimuth_bounds_l = (60-azi, 90-azi) 
        azimuth_bounds_r = (90-azi, 120-azi) 
        altitude_bounds = (65-alt, 80-alt)

    def eval_overlap_from_position(pos, side: str, optimize_size: bool, spine_tensor: torch.Tensor, image_shape, spacing):
        if optimize_size:
            # d, L = snap_to_discrete_values(pos[5], pos[6])
            d, L = snap_to_discrete_values_xfr(pos[5], pos[6])
            params_5 = pos[:5]
        else:
            d, L = diameter, length
            params_5 = pos

        cyl_mask = generate_cylinder_n_torch(
            d, L, params_5[0], params_5[1], params_5[2], params_5[3], params_5[4],
            image_shape, spacing, device, grid
        )
        overlap = compute_overlap_ratio_from_cylinder_mask(cyl_mask, spine_tensor)
        return overlap, d, L

    if optimize_size:
        print("=== DE 最佳化模式：最佳化位置、角度、直徑和長度 ===")
        diameter_bounds = (min(ALLOWED_DIAMETERS), max(ALLOWED_DIAMETERS))
        length_bounds = (min(ALLOWED_LENGTHS), max(ALLOWED_LENGTHS))
        
        bounds_l = [z_bounds, y_bounds, x_bounds_left, azimuth_bounds_l, altitude_bounds, diameter_bounds, length_bounds]
        bounds_r = [z_bounds, y_bounds, x_bounds_right, azimuth_bounds_r, altitude_bounds, diameter_bounds, length_bounds]
    else:
        print("=== DE 固定尺寸模式：最佳化位置和角度 ===")
        diameter = 4.5  
        length = 45     
        
        bounds_l = [z_bounds, y_bounds, x_bounds_left, azimuth_bounds_l, altitude_bounds]
        bounds_r = [z_bounds, y_bounds, x_bounds_right, azimuth_bounds_r, altitude_bounds]

    # DE 的 popsize 實際粒子數 = popsize * len(bounds)
    # 為了跟 PSO 公平比較，我們讓它轉換一下
    de_popsize = max(1, swarm_size // len(bounds_l))

    # --- 左側最佳化 ---
    print("\n=== 左側 (DE) ===")
    res_l = differential_evolution(objective_fn, bounds_l, popsize=de_popsize, maxiter=max_iter)
    position_l, loss_l = res_l.x, res_l.fun

    overlap_l, diameter_l, length_l = eval_overlap_from_position(position_l, "L", optimize_size, spine_tensor, image_shape, spacing)
    best_position_l = list(position_l[:5]) + [diameter_l, length_l] if optimize_size else list(position_l)
    best_loss_l, best_overlap_l = loss_l, overlap_l
    """
    retries = 0
    while (best_loss_l > 0 or best_overlap_l < OVERLAP_THRESH) and retries < 10:
        res_l = differential_evolution(objective_fn, bounds_l, popsize=de_popsize, maxiter=max_iter)
        position_l, loss_l = res_l.x, res_l.fun
        overlap_l, diameter_l, length_l = eval_overlap_from_position(position_l, "L", optimize_size, spine_tensor, image_shape, spacing)
        
        candidate_pos = (list(position_l[:5]) + [diameter_l, length_l]) if optimize_size else list(position_l)
        if is_solution_ok(loss_l, overlap_l, OVERLAP_THRESH) and (not is_solution_ok(best_loss_l, best_overlap_l, OVERLAP_THRESH) or loss_l < best_loss_l):
            best_position_l, best_loss_l, best_overlap_l = candidate_pos, loss_l, overlap_l
        elif (not is_solution_ok(best_loss_l, best_overlap_l, OVERLAP_THRESH)) and (loss_l < best_loss_l):
            best_position_l, best_loss_l, best_overlap_l = candidate_pos, loss_l, overlap_l
        retries += 1
    """
    # --- 右側最佳化 ---
    print("\n=== 右側 (DE) ===")
    res_r = differential_evolution(objective_fn, bounds_r, popsize=de_popsize, maxiter=max_iter)
    position_r, loss_r = res_r.x, res_r.fun

    overlap_r, diameter_r, length_r = eval_overlap_from_position(position_r, "R", optimize_size, spine_tensor, image_shape, spacing)
    best_position_r = list(position_r[:5]) + [diameter_r, length_r] if optimize_size else list(position_r)
    best_loss_r, best_overlap_r = loss_r, overlap_r
    """
    retries = 0
    while (best_loss_r > 0 or best_overlap_r < OVERLAP_THRESH) and retries < 10:
        res_r = differential_evolution(objective_fn, bounds_r, popsize=de_popsize, maxiter=max_iter)
        position_r, loss_r = res_r.x, res_r.fun
        overlap_r, diameter_r, length_r = eval_overlap_from_position(position_r, "R", optimize_size, spine_tensor, image_shape, spacing)
        
        candidate_pos = (list(position_r[:5]) + [diameter_r, length_r]) if optimize_size else list(position_r)
        if is_solution_ok(loss_r, overlap_r, OVERLAP_THRESH) and (not is_solution_ok(best_loss_r, best_overlap_r, OVERLAP_THRESH) or loss_r < best_loss_r):
            best_position_r, best_loss_r, best_overlap_r = candidate_pos, loss_r, overlap_r
        elif (not is_solution_ok(best_loss_r, best_overlap_r, OVERLAP_THRESH)) and (loss_r < best_loss_r):
            best_position_r, best_loss_r, best_overlap_r = candidate_pos, loss_r, overlap_r
        retries += 1
    """
    total_time = time.time() - start_time
    
    final_diameter_l = best_position_l[5] if optimize_size else diameter
    final_length_l = best_position_l[6] if optimize_size else length
    final_diameter_r = best_position_r[5] if optimize_size else diameter
    final_length_r = best_position_r[6] if optimize_size else length

    render_bone_figure(
        None, None,
        spine_tensor, cortical_tensor,
        'Output',
        spacing=spacing,
        way='CBT' if CBT else 'TPS',
        best_position_l=best_position_l, best_position_r=best_position_r,
        diameter_l=final_diameter_l, length_l=final_length_l,
        diameter_r=final_diameter_r, length_r=final_length_r,
        image2_path=image2_path,
        device=device, grid=grid,
        swarm_size=swarm_size, max_iter=max_iter, total_time=total_time,
    )

    return best_position_l, best_loss_l, best_position_r, best_loss_r, total_time

def run_nm_torch(
    label_str: str,
    image1_path: str,
    image2_path: str,
    image3_path: str,
    folder: str,
    swarm_size: int, # NM 不用 swarm_size，但保留參數以維持介面統一
    max_iter: int,
    spacing: list,
    CBT: bool,
    device: torch.device,
    optimize_size: bool = True,
    grid=None
):
    """
    使用 Nelder-Mead 進行最佳化
    """
    start_time = time.time()

    global image1_array, image2_array, image2_shape, image3_array
    global diameter, length  
    global spine_tensor, cortical_tensor, spine_roi_tensor

    image1 = sitk.ReadImage(image1_path)
    image2 = sitk.ReadImage(image2_path)
    image3 = sitk.ReadImage(image3_path)
    image1_array = sitk.GetArrayFromImage(image1)
    image2_array = sitk.GetArrayFromImage(image2)
    image3_array = sitk.GetArrayFromImage(image3)
    image2_shape = image2_array.shape
    image_shape = image2_shape

    cortical_tensor = torch.from_numpy(image1_array).to(device=device, dtype=torch.uint8)
    spine_tensor = torch.from_numpy(image2_array).to(device=device, dtype=torch.uint8)
    spine_roi_tensor = torch.from_numpy(image3_array).to(device=device, dtype=torch.uint8)
    # 建立明確的優化上下文（取代舊的跨檔案 global 注入；狀態全部顯式傳入目標函數）
    ctx = OptimizationContext(
        cortical_tensor=cortical_tensor,
        spine_tensor=spine_tensor,
        spine_roi_tensor=spine_roi_tensor,
        image1_array=image1_array,
        image2_array=image2_array,
        image3_array=image3_array,
        image2_shape=image2_shape,
        spacing=spacing,
        device=device,
        grid=grid,
        diameter=diameter if not optimize_size else None,
        length=length if not optimize_size else None,
    )
    objective_fn = make_objective_function(ctx)

    azi = azimuth_rotation(image2_path)
    res = analyze_vertebral_tilt_contour(image2_path, edge_type='superior', show_plot=False, debug=False)
    alt = res['superior']['tilt_angle_deg']

    if CBT == True:    
        z_bounds = (0, image_shape[0] - 1)
        y_bounds = (image_shape[1]/5, image_shape[1]/2 - 1)
        x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1)
        x_bounds_left  = (0, image_shape[2]/2 - image_shape[2]/10 - 1)
        # CBT 參數依據 references/：Santoni 2009 (Spine J 9:366) 冠狀面 25-30° cranial (caudo-cephalad)、
        # 軸狀面自正中線向外 (medial→lateral) ≤30°；Delgado-Fernandez 2017 (ASJ 11:817)、Kim 2022 (SSRR 6:1)
        azimuth_bounds_l = ((98-azi), (120-azi))
        azimuth_bounds_r = ((60-azi), (82-azi))
        altitude_bounds = ((60-alt), (70-alt))
    else:      
        z_bounds = (0, image_shape[0] - 1)   
        y_bounds = (image_shape[1]/5, image_shape[1]/2 - 1)   
        x_bounds_left = (0, image_shape[2]/2 - image_shape[2]/10 - 1)   
        x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1)
        azimuth_bounds_l = (60-azi, 90-azi) 
        azimuth_bounds_r = (90-azi, 120-azi) 
        altitude_bounds = (65-alt, 80-alt)

    def eval_overlap_from_position(pos, side: str, optimize_size: bool, spine_tensor: torch.Tensor, image_shape, spacing):
        if optimize_size:
            # d, L = snap_to_discrete_values(pos[5], pos[6])
            d, L = snap_to_discrete_values_xfr(pos[5], pos[6])
            params_5 = pos[:5]
        else:
            d, L = diameter, length
            params_5 = pos

        cyl_mask = generate_cylinder_n_torch(
            d, L, params_5[0], params_5[1], params_5[2], params_5[3], params_5[4],
            image_shape, spacing, device, grid
        )
        overlap = compute_overlap_ratio_from_cylinder_mask(cyl_mask, spine_tensor)
        return overlap, d, L

    if optimize_size:
        print("=== NM 最佳化模式 ===")
        bounds_l = [z_bounds, y_bounds, x_bounds_left, azimuth_bounds_l, altitude_bounds, 
                    (min(ALLOWED_DIAMETERS), max(ALLOWED_DIAMETERS)), (min(ALLOWED_LENGTHS), max(ALLOWED_LENGTHS))]
        bounds_r = [z_bounds, y_bounds, x_bounds_right, azimuth_bounds_r, altitude_bounds, 
                    (min(ALLOWED_DIAMETERS), max(ALLOWED_DIAMETERS)), (min(ALLOWED_LENGTHS), max(ALLOWED_LENGTHS))]
    else:
        print("=== NM 固定尺寸模式 ===")
        diameter, length = 4.5, 45     
        bounds_l = [z_bounds, y_bounds, x_bounds_left, azimuth_bounds_l, altitude_bounds]
        bounds_r = [z_bounds, y_bounds, x_bounds_right, azimuth_bounds_r, altitude_bounds]

    def get_random_x0(bounds):
        # 產生在 Bounds 內的隨機起始點
        return [np.random.uniform(b[0], b[1]) for b in bounds]

    # --- 左側最佳化 ---
    print("\n=== 左側 (Nelder-Mead) ===")
    x0_l = get_random_x0(bounds_l)
    res_l = minimize(objective_fn, x0_l, method='Nelder-Mead', bounds=bounds_l, options={'maxiter': max_iter})
    position_l, loss_l = res_l.x, res_l.fun

    overlap_l, diameter_l, length_l = eval_overlap_from_position(position_l, "L", optimize_size, spine_tensor, image_shape, spacing)
    best_position_l = list(position_l[:5]) + [diameter_l, length_l] if optimize_size else list(position_l)
    best_loss_l, best_overlap_l = loss_l, overlap_l
    
    retries = 0
    while (best_loss_l > 0 or best_overlap_l < OVERLAP_THRESH) and retries < 10:
        x0_l = get_random_x0(bounds_l) # 每次 retry 都換一個隨機起始點
        res_l = minimize(objective_fn, x0_l, method='Nelder-Mead', bounds=bounds_l, options={'maxiter': max_iter})
        position_l, loss_l = res_l.x, res_l.fun
        overlap_l, diameter_l, length_l = eval_overlap_from_position(position_l, "L", optimize_size, spine_tensor, image_shape, spacing)
        
        candidate_pos = (list(position_l[:5]) + [diameter_l, length_l]) if optimize_size else list(position_l)
        if is_solution_ok(loss_l, overlap_l, OVERLAP_THRESH) and (not is_solution_ok(best_loss_l, best_overlap_l, OVERLAP_THRESH) or loss_l < best_loss_l):
            best_position_l, best_loss_l, best_overlap_l = candidate_pos, loss_l, overlap_l
        elif (not is_solution_ok(best_loss_l, best_overlap_l, OVERLAP_THRESH)) and (loss_l < best_loss_l):
            best_position_l, best_loss_l, best_overlap_l = candidate_pos, loss_l, overlap_l
        retries += 1

    # --- 右側最佳化 ---
    print("\n=== 右側 (Nelder-Mead) ===")
    x0_r = get_random_x0(bounds_r)
    res_r = minimize(objective_fn, x0_r, method='Nelder-Mead', bounds=bounds_r, options={'maxiter': max_iter})
    position_r, loss_r = res_r.x, res_r.fun

    overlap_r, diameter_r, length_r = eval_overlap_from_position(position_r, "R", optimize_size, spine_tensor, image_shape, spacing)
    best_position_r = list(position_r[:5]) + [diameter_r, length_r] if optimize_size else list(position_r)
    best_loss_r, best_overlap_r = loss_r, overlap_r

    retries = 0
    while (best_loss_r > 0 or best_overlap_r < OVERLAP_THRESH) and retries < 10:
        x0_r = get_random_x0(bounds_r)
        res_r = minimize(objective_fn, x0_r, method='Nelder-Mead', bounds=bounds_r, options={'maxiter': max_iter})
        position_r, loss_r = res_r.x, res_r.fun
        overlap_r, diameter_r, length_r = eval_overlap_from_position(position_r, "R", optimize_size, spine_tensor, image_shape, spacing)
        
        candidate_pos = (list(position_r[:5]) + [diameter_r, length_r]) if optimize_size else list(position_r)
        if is_solution_ok(loss_r, overlap_r, OVERLAP_THRESH) and (not is_solution_ok(best_loss_r, best_overlap_r, OVERLAP_THRESH) or loss_r < best_loss_r):
            best_position_r, best_loss_r, best_overlap_r = candidate_pos, loss_r, overlap_r
        elif (not is_solution_ok(best_loss_r, best_overlap_r, OVERLAP_THRESH)) and (loss_r < best_loss_r):
            best_position_r, best_loss_r, best_overlap_r = candidate_pos, loss_r, overlap_r
        retries += 1

    total_time = time.time() - start_time
    
    final_diameter_l = best_position_l[5] if optimize_size else diameter
    final_length_l = best_position_l[6] if optimize_size else length
    final_diameter_r = best_position_r[5] if optimize_size else diameter
    final_length_r = best_position_r[6] if optimize_size else length

    render_bone_figure(
        None, None,
        spine_tensor, cortical_tensor,
        'Output',
        spacing=spacing,
        way='CBT' if CBT else 'TPS',
        best_position_l=best_position_l, best_position_r=best_position_r,
        diameter_l=final_diameter_l, length_l=final_length_l,
        diameter_r=final_diameter_r, length_r=final_length_r,
        image2_path=image2_path,
        device=device, grid=grid,
        swarm_size=swarm_size, max_iter=max_iter, total_time=total_time,
    )

    return best_position_l, best_loss_l, best_position_r, best_loss_r, total_time