import os, slicer, shutil, SimpleITK as sitk

def postprocess(root_dir):
    os.makedirs(root_dir + "\\nii_2_post", exist_ok=True)
    os.makedirs(root_dir + "\\stl_1", exist_ok=True)

    listdir = os.listdir(root_dir + "\\nii_1\\")
    listdir.sort()
    for i in range(len(listdir)):
        if not listdir[i].endswith(".nii.gz"): continue
        if os.path.exists(root_dir + "\\stl_1\\" + listdir[i].replace(".nii.gz", ".stl")): continue 
        
        print(f"{i}/{len(listdir)}, {listdir[i]}")

        sitk_image = sitk.ReadImage(root_dir + "\\nii_1\\" + listdir[i])
        if sitk.GetArrayFromImage(sitk_image).max() == 0:
            continue

        masterVolumeNode = slicer.util.loadVolume(root_dir + "\\nii_1\\" + listdir[i])

        # Create segmentation
        segmentationNode = slicer.util.loadSegmentation(root_dir + "\\nii_1\\" + listdir[i])

        # Create segment editor to get access to effects
        segmentEditorWidget = slicer.qMRMLSegmentEditorWidget()
        segmentEditorWidget.setMRMLScene(slicer.mrmlScene)
        segmentEditorNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentEditorNode")
        segmentEditorWidget.setMRMLSegmentEditorNode(segmentEditorNode)
        segmentEditorWidget.setSegmentationNode(segmentationNode)
        segmentEditorWidget.setMasterVolumeNode(masterVolumeNode)

        # Smoothing
        segmentEditorWidget.setActiveEffectByName("Smoothing")
        effect = segmentEditorWidget.activeEffect()
        effect.setParameter("SmoothingMethod", "MEDIAN")
        effect.setParameter("KernelSizeMm", 1)
        effect.self().onApply()

        # Systematically remove small islands
        segmentEditorWidget.setActiveEffectByName("Islands")
        effect = segmentEditorWidget.activeEffect()
        effect.setParameter("MinimumSize","1000")
        effect.setParameter("Operation","KEEP_LARGEST_ISLAND")
        segmentEditorNode.SetOverwriteMode(slicer.vtkMRMLSegmentEditorNode.OverwriteNone) 
        effect.self().onApply()

        labelmapVolumeNode = slicer.mrmlScene.AddNewNodeByClass('vtkMRMLLabelMapVolumeNode')
        slicer.modules.segmentations.logic().ExportVisibleSegmentsToLabelmapNode(segmentationNode, labelmapVolumeNode, masterVolumeNode)
        slicer.util.saveNode(labelmapVolumeNode, root_dir + "\\nii_2_post\\" + listdir[i])

        # Clean up
        segmentEditorWidget = None
        slicer.mrmlScene.RemoveNode(segmentEditorNode)

        # Make segmentation results visible in 3D
        segmentationNode.CreateClosedSurfaceRepresentation()

        # Write to STL file
        slicer.vtkSlicerSegmentationsModuleLogic.ExportSegmentsClosedSurfaceRepresentationToFiles(root_dir + "\\stl_1\\", segmentationNode, None, "STL")
        slicer.mrmlScene.Clear(0)

        shutil.move(root_dir + "\\stl_1\\" + listdir[i] + "_Segment_1.stl", 
                    root_dir + "\\stl_1\\" + listdir[i].replace(".nii.gz", ".stl"))

if __name__ == "__main__":
    preprocess(r"preprocess\\")