import os
import SimpleITK as sitk
from config.constant import LABEL_MAP
from imaging.resample import resample_img
import numpy as np

"""
# 沿用原本 LABEL_MAP
seg_bone(n, name, img, lbl)

# user 自定義
my_map = {1: "L1", 2: "L2", 3: "L3"}
seg_bone(n, name, img, lbl, label_map=my_map)
"""

def _largest_cc_bbox(mask_img):
    """26-連通的最大连通區域 + 其 bbox（RelabelComponent 依大小排序，最大者=1）。
    回傳 (largest_mask, bbox2)；沒有任何組件時回傳 None。
    bbox2 格式：[x_start, y_start, z_start, x_size, y_size, z_size]。"""
    cc_image = sitk.ConnectedComponent(mask_img, True)  # fullyConnected
    relabeled_cc = sitk.RelabelComponent(cc_image, sortByObjectSize=True)
    shape_stats = sitk.LabelShapeStatisticsImageFilter()
    shape_stats.Execute(relabeled_cc)
    if shape_stats.GetNumberOfLabels() < 1:
        return None
    return (relabeled_cc == 1), shape_stats.GetBoundingBox(1)

def _bbox_roi(img, bbox, margin=0):
    """裁 bbox（對稱外擴 margin 個 voxel，clamp 到影像邊界）。
    bbox 格式：[x_start, y_start, z_start, x_size, y_size, z_size]。"""
    n = img.GetSize()  # (x, y, z)
    index = [max(0, int(bbox[i]) - margin) for i in range(3)]
    size = [min(n[i] - index[i], int(bbox[i + 3]) + 2 * margin) for i in range(3)]
    return sitk.RegionOfInterest(img, size, index)

def seg_bone(n, name, resampled_sitk_img, resampled_sitk_lbl, output_base=None, label_map=LABEL_MAP,
             original_label=None):

    if output_base==None:
        output_base=='Dataset'

    if n not in label_map:
        raise ValueError(f"Label {n} not found in label_map")

    label_name = label_map[n]

    # ============ 原解析度（未插值）chain ============
    # 1. 提取標籤 n 的二值遮罩 (將標籤 n 設為 1，其餘為 0)，最大連通區域
    smd_path = resampled_path = binary_sdf_path = binary_erode_path = None
    binary_linear_path = binary_nn_path = None
    if original_label is not None:
        bin_orig = sitk.BinaryThreshold(original_label, n, n, 1, 0)
        cc_orig = _largest_cc_bbox(bin_orig)
        if cc_orig is None:
            return None
        largest_orig, bbox_orig = cc_orig

        # _binary.nii.gz：原解析度【未插值】遮罩（最大连通區域、裁到物件 bbox，
        # 不重取樣、不插值 —— 原始 label 的忠實二值版本）
        binary_path = os.path.join(output_base, f"{label_name}_binary.nii.gz")
        sitk.WriteImage(_bbox_roi(largest_orig, bbox_orig), binary_path)

        # _smd.nii.gz：SignedMaurerDistanceMap（ITK 慣例：物件內負 / 外正；
        # 距離以原始 index（pixel）單位、不隨各向异性 spacing 縮放——
        # _binary_sdf 的 0.5 閾值（mid-gap）正是依賴這個 index 單位慣例）。
        # 裁 bbox_orig 外擴 4 voxel 的背景輪：
        # 沒有背景輪時填充值直接貼着物件邊緣，重取樣會在裁切邊界產生
        # 假的閾值穿越。
        # 這個 SimpleITK build 的 3D SignedMaurerDistanceMap 只支援整數輸入，
        # 先 Cast 到 uint8 再回傳 float32 輸出
        smd_full = sitk.SignedMaurerDistanceMap(sitk.Cast(largest_orig, sitk.sitkUInt8))
        smd_full = sitk.Cast(smd_full, sitk.sitkFloat32)
        smd_margined = _bbox_roi(smd_full, bbox_orig, margin=4)
        smd_path = os.path.join(output_base, f"{label_name}_smd.nii.gz")
        sitk.WriteImage(smd_margined, smd_path)

        # SMD 體積【线性插值】重取樣到 0.5mm（reference = 0.5mm CT，與
        # resampled_sitk_img 同 grid，之後可直接用 0.5mm bbox 裁切）。
        # SMD 在每個 input voxel 內分段線性，线性重取樣近似精確、無漣波，
        # 各等值面（物件邊界等）不變。填充值 = 裁切角落（背景側，正值）；
        # 若為負（物件貼影像邊界的病態情況）用 0。
        corner = float(sitk.GetArrayViewFromImage(smd_margined).flat[0])
        rs = sitk.ResampleImageFilter()
        rs.SetReferenceImage(resampled_sitk_img)
        rs.SetInterpolator(sitk.sitkLinear)
        rs.SetDefaultPixelValue(corner if corner > 0 else 0.0)
        smd_res_full = rs.Execute(smd_margined)

        # 0.5mm linear 二值化 mask（full extent，與 resampled_sitk_img 同 grid）：
        # 原數據 5mm 切片上採樣 10x 到 0.5mm，最近邻會在邊界產生 10 體素厚的
        # 階梯鋸齒；线性插值使邊界落在次體素位置（rotation 前的邊界更平滑）。
        bin_lin = resample_img(sitk.Cast(bin_orig, sitk.sitkFloat32))
        arr = (sitk.GetArrayFromImage(bin_lin) > 0.5).astype(np.uint8)
        if arr.shape != resampled_sitk_img.GetSize()[::-1]:
            raise RuntimeError(
                f"linear binary resample shape mismatch: {arr.shape} vs {resampled_sitk_img.GetSize()}")
        binary_mask = sitk.GetImageFromArray(arr)
        binary_mask.CopyInformation(resampled_sitk_img)
    else:
        # 無原始 label：舊版路徑（0.5mm label 閾值），不產生 SMD/SDF chain
        binary_mask = sitk.BinaryThreshold(resampled_sitk_lbl, n, n, 1, 0)
        binary_path = None

    # 2. 0.5mm 最大連通區域（26-連通）+ 邊界框；所有 0.5mm 輸出裁到同一 bbox
    cc_res = _largest_cc_bbox(binary_mask)
    if cc_res is None:
        return None
    largest_mask, bbox2 = cc_res

    if binary_path is None:
        binary_path = os.path.join(output_base, f"{label_name}_binary.nii.gz")
        sitk.WriteImage(sitk.RegionOfInterest(largest_mask, bbox2[3:], bbox2[:3]), binary_path)

    if smd_path is not None:
        # _smd_resampled.nii.gz：线性重取樣到 0.5mm 的 SMD（浮點，裁 bbox2）
        resampled_path = os.path.join(output_base, f"{label_name}_smd_resampled.nii.gz")
        sitk.WriteImage(sitk.RegionOfInterest(smd_res_full, bbox2[3:], bbox2[:3]), resampled_path)

        # _binary_sdf.nii.gz：_smd_resampled 於 0.5 閾值 -> 0.5mm 平滑 mask
        # （SMD 內負/外正；0.5 介於內殼 ≈0 與外殼 ≈+1 之間，即原解析度
        #  label 邊界的 mid-gap 位置，physical volume 與 _binary/_binary_nn
        #  一致，sub-voxel 表面、無 NN 階梯）
        bin_sdf_full = sitk.GetImageFromArray(
            (sitk.GetArrayFromImage(smd_res_full) < 0.5).astype(np.uint8))
        bin_sdf_full.CopyInformation(resampled_sitk_img)
        binary_sdf_path = os.path.join(output_base, f"{label_name}_binary_sdf.nii.gz")
        sitk.WriteImage(sitk.RegionOfInterest(bin_sdf_full, bbox2[3:], bbox2[:3]), binary_sdf_path)

        # _binary_nn.nii.gz：0.5mm label 最近邻（舊版），裁自己的 bbox，僅供對比
        nn_mask = sitk.BinaryThreshold(resampled_sitk_lbl, n, n, 1, 0)
        nn_res = _largest_cc_bbox(nn_mask)
        if nn_res is not None:
            nn_largest, nn_bbox = nn_res
            binary_nn_path = os.path.join(output_base, f"{label_name}_binary_nn.nii.gz")
            sitk.WriteImage(sitk.RegionOfInterest(nn_largest, nn_bbox[3:], nn_bbox[:3]),
                            binary_nn_path)

    # 3. roi（0.5mm）
    # _roi2 不再存檔；_cortical 改由 xfr_preprocess 的旋轉後處理產出
    # （rotated/{level}_cortical.nii.gz，定義不變：門檻 = 骨頭 mask 內 median HU）
    roi    = sitk.RegionOfInterest(resampled_sitk_img, bbox2[3:], bbox2[:3])
    roi_path = os.path.join(output_base, f"{label_name}_roi.nii.gz")
    sitk.WriteImage(roi, roi_path)

    return roi_path, binary_path, None, None, binary_nn_path, \
        binary_linear_path, smd_path, resampled_path, binary_sdf_path, binary_erode_path

"""
Dataset/
└── standardized/
    └── subject001/
        ├── L1_binary.nii.gz          # 原解析度【未插值】遮罩（最大连通區域、裁物件 bbox）
        ├── L1_smd.nii.gz             # SignedMaurerDistanceMap（內負/外正，原始 index 單位；
        │                             #   bbox 外扩 4 voxel 背景輪，供重取樣插值用）
        ├── L1_smd_resampled.nii.gz   # _smd 經线性插值重取樣到 0.5mm（浮點，裁 0.5mm bbox）
        ├── L1_binary_sdf.nii.gz      # _smd_resampled 於 0.5 閾值 -> 0.5mm 平滑 mask（0/1，裁同 bbox）
        ├── L1_binary_nn.nii.gz       # 最近邻版 0/1（對比用，各自 bbox）
        ├── L1_roi.nii.gz
        ├── L2_binary.nii.gz
        ...
"""