import SimpleITK as sitk
from rt_utils import RTStructBuilder

import numpy as np

DCM_CT = '/nn/3378417/20241112/CT/8'
NII_CT = "/nn/3378417/20241112/nii/8_1.1_CyberKnife_head(MAR)_20241112110016_4.nii.gz"
label_file = "/nn/3378417/20241112/nii/8_1.1_CyberKnife_head(MAR)_20241112110016.label.nii.gz"
rtss_file = '/nn/3378417/20241112/CT/8-rtss.dcm'

# reg_transform(NII_CT, NII_MR, output_file, label_file)

reader = sitk.ImageSeriesReader()
dicom_names = reader.GetGDCMSeriesFileNames(DCM_CT)
reader.SetFileNames(dicom_names)
# reader.SetFileNames(sorted(dicom_names))
reader.MetaDataDictionaryArrayUpdateOn()
reader.LoadPrivateTagsOn()
image = reader.Execute()
sitk.WriteImage(image, '0/image.nii.gz')

nnU = sitk.ReadImage(label_file)
nnU = sitk.Resample(nnU, image, sitk.Transform(), sitk.sitkNearestNeighbor)

sitk.WriteImage(nnU, '0/nnU.nii.gz')

ccfilter  = sitk.ConnectedComponentImageFilter ()
nnUCC = ccfilter.Execute(nnU)
ObjectCount1 = ccfilter.GetObjectCount()

rtstruct = RTStructBuilder.create_new(dicom_series_path=DCM_CT)

print(ObjectCount1)
for j1 in range(ObjectCount1):
    label1  = sitk.BinaryThreshold(nnUCC, j1+1, j1+1)
    # label1  = sitk.AntiAliasBinary(label1)

    mask = sitk.GetArrayFromImage(label1).astype(bool)
    mask = np.transpose(mask, (1, 2, 0))
    # continue
    if mask.any():
        print(j1)
        rtstruct.add_roi(
            mask=mask,
            # use_pin_hole=True,
            # name="n%d"%n,
            )

print(rtss_file)
rtstruct.save(rtss_file)

# return rtss_file