import os
import SimpleITK as sitk
from config.constant import LABEL_MAP
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 seg_bone(n, name, resampled_sitk_img, resampled_sitk_lbl, output_base=None, label_map=LABEL_MAP):

    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]

    lssif = sitk.LabelShapeStatisticsImageFilter()
    lssif.Execute(resampled_sitk_lbl)

    if not lssif.HasLabel(n):
        raise RuntimeError(f"Label {n} not found")

    bbox2 = lssif.GetBoundingBox(n)

    roi    = sitk.RegionOfInterest(resampled_sitk_img, bbox2[3:], bbox2[:3])
    label2 = sitk.RegionOfInterest(resampled_sitk_lbl, bbox2[3:], bbox2[:3])
    roi_path = os.path.join(output_base, f"{label_name}_roi.nii.gz")
    sitk.WriteImage(roi, roi_path)

    binary = sitk.BinaryThreshold(label2, lowerThreshold=n, upperThreshold=n, outsideValue=0, insideValue=1)
    binary_path = os.path.join(output_base, f"{label_name}_binary.nii.gz")
    sitk.WriteImage(binary, binary_path)

    roi_pixel_type = roi.GetPixelID()
    binary_cast = sitk.Cast(binary, roi_pixel_type)
    roi2 = roi * binary_cast
    roi2_path = os.path.join(output_base, f"{label_name}_roi2.nii.gz")
    sitk.WriteImage(roi2, roi2_path)

    lsif = sitk.LabelStatisticsImageFilter()
    label2_int = sitk.Cast(label2, sitk.sitkUInt16)
    lsif.Execute(roi2, label2_int)
    labels_in_roi = lsif.GetLabels()
    if n in labels_in_roi:
        roi_hu = sitk.GetArrayFromImage(roi2)
        threshold = np.percentile(roi_hu, 60)
    else:
        threshold = lsif.GetMedian(labels_in_roi[0])

    cortical = sitk.BinaryThreshold(roi2, lowerThreshold=threshold, upperThreshold=10000, outsideValue=0, insideValue=1)
    cortical_path = os.path.join(output_base, f"{label_name}_cortical.nii.gz")
    sitk.WriteImage(cortical, cortical_path)

    return roi_path, binary_path, roi2_path, cortical_path

"""
Dataset/
└── standardized/
    └── subject001/
        ├── L1_roi.nii.gz
        ├── L1_binary.nii.gz
        ├── L1_roi2.nii.gz
        ├── L1_cortical.nii.gz
        ├── L2_roi.nii.gz
        ├── L2_binary.nii.gz
        ...
"""