import logging
import os
import numpy as np
import vtk
import ctk
import qt
import slicer
from slicer.ScriptedLoadableModule import *
from slicer.util import VTKObservationMixin
import SimpleITK as sitk
import sitkUtils



#
# HeartSegmentation
#

class HeartSegmentation(ScriptedLoadableModule):
    """Uses ScriptedLoadableModule base class, available at:
    https://github.com/Slicer/Slicer/blob/main/Base/Python/slicer/ScriptedLoadableModule.py
    """

    def __init__(self, parent):
        ScriptedLoadableModule.__init__(self, parent)
        self.parent.title = "HeartSegmentation"  # TODO: make this more human readable by adding spaces
        self.parent.categories = ["Examples"]  # TODO: set categories (folders where the module shows up in the module selector)
        self.parent.dependencies = []  # TODO: add here list of module names that this module requires
        self.parent.contributors = ["John Doe (AnyWare Corp.)"]  # TODO: replace with "Firstname Lastname (Organization)"
        # TODO: update with short description of the module and a link to online module documentation
        self.parent.helpText = """
This is an example of scripted loadable module bundled in an extension.
See more information in <a href="https://github.com/organization/projectname#HeartSegmentation">module documentation</a>.
"""
        # TODO: replace with organization, grant and thanks
        self.parent.acknowledgementText = """
This file was originally developed by Jean-Christophe Fillion-Robin, Kitware Inc., Andras Lasso, PerkLab,
and Steve Pieper, Isomics, Inc. and was partially funded by NIH grant 3P41RR013218-12S1.
"""

        # Additional initialization step after application startup is complete
        slicer.app.connect("startupCompleted()", registerSampleData)


#
# Register sample data sets in Sample Data module
#

def registerSampleData():
    """
    Add data sets to Sample Data module.
    """
    # It is always recommended to provide sample data for users to make it easy to try the module,
    # but if no sample data is available then this method (and associated startupCompeted signal connection) can be removed.

    import SampleData
    iconsPath = os.path.join(os.path.dirname(__file__), 'Resources/Icons')

    # To ensure that the source code repository remains small (can be downloaded and installed quickly)
    # it is recommended to store data sets that are larger than a few MB in a Github release.

    # HeartSegmentation1
    SampleData.SampleDataLogic.registerCustomSampleDataSource(
        # Category and sample name displayed in Sample Data module
        category='HeartSegmentation',
        sampleName='HeartSegmentation1',
        # Thumbnail should have size of approximately 260x280 pixels and stored in Resources/Icons folder.
        # It can be created by Screen Capture module, "Capture all views" option enabled, "Number of images" set to "Single".
        thumbnailFileName=os.path.join(iconsPath, 'HeartSegmentation1.png'),
        # Download URL and target file name
        uris="https://github.com/Slicer/SlicerTestingData/releases/download/SHA256/998cb522173839c78657f4bc0ea907cea09fd04e44601f17c82ea27927937b95",
        fileNames='HeartSegmentation1.nrrd',
        # Checksum to ensure file integrity. Can be computed by this command:
        #  import hashlib; print(hashlib.sha256(open(filename, "rb").read()).hexdigest())
        checksums='SHA256:998cb522173839c78657f4bc0ea907cea09fd04e44601f17c82ea27927937b95',
        # This node name will be used when the data set is loaded
        nodeNames='HeartSegmentation1'
    )

    # HeartSegmentation2
    SampleData.SampleDataLogic.registerCustomSampleDataSource(
        # Category and sample name displayed in Sample Data module
        category='HeartSegmentation',
        sampleName='HeartSegmentation2',
        thumbnailFileName=os.path.join(iconsPath, 'HeartSegmentation2.png'),
        # Download URL and target file name
        uris="https://github.com/Slicer/SlicerTestingData/releases/download/SHA256/1a64f3f422eb3d1c9b093d1a18da354b13bcf307907c66317e2463ee530b7a97",
        fileNames='HeartSegmentation2.nrrd',
        checksums='SHA256:1a64f3f422eb3d1c9b093d1a18da354b13bcf307907c66317e2463ee530b7a97',
        # This node name will be used when the data set is loaded
        nodeNames='HeartSegmentation2'
    )


#
# HeartSegmentationWidget
#

class HeartSegmentationWidget(ScriptedLoadableModuleWidget, VTKObservationMixin):
    """Uses ScriptedLoadableModuleWidget base class, available at:
    https://github.com/Slicer/Slicer/blob/main/Base/Python/slicer/ScriptedLoadableModule.py
    """
    def __init__(self, parent=None):
        """
        Called when the user opens the module the first time and the widget is initialized.
        """
        ScriptedLoadableModuleWidget.__init__(self, parent)
        VTKObservationMixin.__init__(self)  # needed for parameter node observation
        self.logic = None
        self._parameterNode = None
        self._updatingGUIFromParameterNode = False
        self.parameterSetNode = None
        self.activeEffets = ["Paint", "Scissors","Erase","Islands","Margin","Smoothing"]

    def setup(self):
        """
        Called when the user opens the module the first time and the widget is initialized.
        """
        ScriptedLoadableModuleWidget.setup(self)

        # Load widget from .ui file (created by Qt Designer).
        uiWidget = slicer.util.loadUI(self.resourcePath('UI/HeartSemgnetation.ui'))
        self.layout.addWidget(uiWidget)
        self.ui = slicer.util.childWidgetVariables(uiWidget)

        # Set scene in MRML widgets. Make sure that in Qt designer the top-level qMRMLWidget's
        # "mrmlSceneChanged(vtkMRMLScene*)" signal in is connected to each MRML widget's.
        # "setMRMLScene(vtkMRMLScene*)" slot.
        uiWidget.setMRMLScene(slicer.mrmlScene)

        # Logic class. Logic implements all computations that should be possible to run
        self.logic = HeartSegmentationLogic()

        # Connections

        self.addObserver(slicer.mrmlScene, slicer.mrmlScene.StartCloseEvent, self.onSceneStartClose)
        self.addObserver(slicer.mrmlScene, slicer.mrmlScene.EndCloseEvent, self.onSceneEndClose)

        self.ui.SegmentsInNode.addItem("None")
        self.ui.ListofFilters.addItems(["RankImageFilter", "MedianImageFilter"])


        self.selectParameterNode()
        self.ui.SegmentEditorWidget.setMRMLScene(slicer.mrmlScene)
        #"vtkMRMLSegmentEditorNode" is Singleton >>


        # These connections ensure that whenever user c.hanges some settings on the GUI, that is saved in the MRML scene
        # (in the selected parameter node).
        #nastavit

        # connections - node related
        self.ui.inputSelector.connect("currentNodeChanged(vtkMRMLNode*)", self.updateParameterNodeFromGUI)
        self.ui.inputSelector.connect("currentNodeChanged(vtkMRMLNode*)",self.displaySelectedVolume)  
        self.ui.outputSelector.connect("currentNodeChanged(vtkMRMLNode*)", self.updateParameterNodeFromGUI)
        self.ui.SegmentationSelector.connect("currentNodeChanged(vtkMRMLNode*)", self.updateParameterNodeFromGUI)
        self.ui.SegmentationSelector.connect("currentNodeChanged(vtkMRMLNode*)", self.refreshSegments)
        self.ui.SegmentsInNode.connect("currentIndexChanged(QString)",self.updateParameterNodeFromGUI)

        self.ui.NodeBoxROI.connect("currentNodeChanged(vtkMRMLNode*)", self.updateParameterNodeFromGUI)

        # connections- spinBox
        self.ui.samplingSpinBox.connect("valueChanged(double)",self.updateParameterNodeFromGUI)
        self.ui.ListofFilters.connect("currentIndexChanged(int)", self.updateParameterNodeFromGUI)

        # connections - checkBox
        self.ui.filtercheckBox.connect("toggled(bool)", self.updateParameterNodeFromGUI)
        self.ui.DownsamplingcheckBox.connect("toggled(bool)", self.updateParameterNodeFromGUI)

        #selection of Segments for Segmentation
        self.ui.SegmentNotspecifiedCheckBox.connect("toggled(bool)", self.SegmentsForSegmentation)
        self.ui.WholeHeartCheckBox.connect("toggled(bool)", self.SegmentsForSegmentation)
        self.ui.LeftVentricleCheckBox.connect("toggled(bool)", self.SegmentsForSegmentation)
        self.ui.RightVentricleCheckBox.connect("toggled(bool)", self.SegmentsForSegmentation)
        self.ui.LeftAtriumCheckBox.connect("toggled(bool)", self.SegmentsForSegmentation)
        self.ui.RightAtriumCheckBox.connect("toggled(bool)", self.SegmentsForSegmentation)

        # connextion - radioButtons
        self.ui.radioButtonNN.connect("toggled(bool)",self.SelectInterpolation)
        self.ui.radioButtonLinear.connect("toggled(bool)",self.SelectInterpolation)
        self.ui.radioButtonWS.connect("toggled(bool)",self.SelectInterpolation)
        self.ui.radioButtonBspline.connect("toggled(bool)",self.SelectInterpolation)

        self.ui.Mask32Button.connect("toggled(bool)",self.SelectSizeOfMask)
        self.ui.Mask64Button.connect("toggled(bool)",self.SelectSizeOfMask)
        self.ui.Mask128Button.connect("toggled(bool)",self.SelectSizeOfMask)
        #self.ui.invertedOutputSelector.connect("currentNodeChanged(vtkMRMLNode*)", self.updateParameterNodeFromGUI)

        # Buttons
        self.ui.ROIVisibleButton.connect('clicked(bool)', self.SetVisibility)
        self.ui.ROIVisibleButton.setIcon(qt.QIcon(self.resourcePath('Icons/view.png')))
        self.ui.MaskSetUpButton.connect('clicked(bool)', self.applyCreatMaskROI)
        self.ui.CropVolumeButton.connect('clicked(bool)', self.apllyCropVolume)
        self.ui.WindowLevelButton.connect('clicked(bool)', self.applyWindowLevel)
        self.ui.CreatROI.connect('clicked(bool)', self.applyCreatROI)
        self.ui.FiltrationButton.connect('clicked(bool)', self.applyFiltration)
        self.ui.StartSegmentationButton.connect('clicked(bool)', self.startSegmentation)
        self.ui.FinishSegmentationButton.connect('clicked(bool)', self.endSegmentation)
        
        self.ui.UpdateButton.connect('clicked(bool)', self.Watershed)
        self.ui.SetBackgroundButton.connect('clicked(bool)', self.setBackground)
        self.ui.applyWaterShedButton.connect('clicked(bool)', self.Watershed)
        self.ui.CancleButton.connect('clicked(bool)', self.Cancel_stopInitalize)
        self.ui.CreatTestDataButton.connect('clicked(bool)', self.creatTestData)

        #hiddent widget
        self.ui.finalSegmentationgroupBox.hide()
        # Make sure parameter node is initialized (needed for module reload)
        self.initializeParameterNode()

    def selectParameterNode(self):
        # Select parameter set node if one is found in the scene, and create one otherwise
        #Contribution: https://github.com/Slicer/Slicer/blob/6931e0d9820a86be8c97f77225bdd3a479980820/Modules/Scripted/SegmentEditor/SegmentEditor.py#L77-L90
        segmentEditorSingletonTag = "SegmentEditor"
        segmentEditorNode = slicer.mrmlScene.GetSingletonNode(segmentEditorSingletonTag, "vtkMRMLSegmentEditorNode")
        if segmentEditorNode is None:
            segmentEditorNode = slicer.mrmlScene.CreateNodeByClass("vtkMRMLSegmentEditorNode")
            segmentEditorNode.UnRegister(None)
            segmentEditorNode.SetSingletonTag(segmentEditorSingletonTag)
            segmentEditorNode = slicer.mrmlScene.AddNode(segmentEditorNode)
        if self.parameterSetNode == segmentEditorNode:
            # nothing changed
                return
        self.parameterSetNode = segmentEditorNode
        self.ui.SegmentEditorWidget.setMRMLSegmentEditorNode(self.parameterSetNode)

    def cleanup(self):
        """
        Called when the application closes and the module widget is destroyed.
        """
        self.removeObservers()

    def enter(self):
        """
        Called each time the user opens this module.
        """
        # Make sure parameter node exists and observed
        self.initializeParameterNode()

    def exit(self):
        """
        Called each time the user opens a different module.
        """
        # Do not react to parameter node changes (GUI wlil be updated when the user enters into the module)
        self.removeObserver(self._parameterNode, vtk.vtkCommand.ModifiedEvent, self.updateGUIFromParameterNode)

    def onSceneStartClose(self, caller, event):
        """
        Called just before the scene is closed.
        """
        # Parameter node will be reset, do not use it anymore
        self.setParameterNode(None)

    def onSceneEndClose(self, caller, event):
        """
        Called just after the scene is closed.
        """
        # If this module is shown while the scene is closed then recreate a new parameter node immediately
        if self.parent.isEntered:
            self.initializeParameterNode()

    def initializeParameterNode(self):
        """
        Ensure parameter node exists and observed.
        """
        # Parameter node stores all user choices in parameter values, node selections, etc.
        # so that when the scene is saved and reloaded, these settings are restored.

        self.setParameterNode(self.logic.getParameterNode())

        # Select default input nodes if nothing is selected yet to save a few clicks for the user
        if not self._parameterNode.GetNodeReference("InputVolume"):
            firstVolumeNode = slicer.mrmlScene.GetFirstNodeByClass("vtkMRMLScalarVolumeNode")
            if firstVolumeNode:
                self._parameterNode.SetNodeReferenceID("InputVolume", firstVolumeNode.GetID())

        # set these parametrs when start>> ensures good behavior when is modul reloaded 
        self._parameterNode.SetParameter("activeSegments", "00000")

        self._parameterNode.SetParameter("cropImageID","None")
 
        self._parameterNode.SetParameter("updateHidden","true")

        self._parameterNode.SetParameter("SegmentationInProcessID","None")


    def setParameterNode(self, inputParameterNode):
        """
        Set and observe parameter node.
        Observation is needed because when the parameter node is changed then the GUI must be updated immediately.
        """

        if inputParameterNode:
            self.logic.setDefaultParameters(inputParameterNode)

        # Unobserve previously selected parameter node and add an observer to the newly selected.
        # Changes of parameter node are observed so that whenever parameters are changed by a script or any other module
        # those are reflected immediately in the GUI.
        if self._parameterNode is not None and self.hasObserver(self._parameterNode, vtk.vtkCommand.ModifiedEvent, self.updateGUIFromParameterNode):
            self.removeObserver(self._parameterNode, vtk.vtkCommand.ModifiedEvent, self.updateGUIFromParameterNode)
        self._parameterNode = inputParameterNode
        if self._parameterNode is not None:
            self.addObserver(self._parameterNode, vtk.vtkCommand.ModifiedEvent, self.updateGUIFromParameterNode)
        # Initial GUI update
        self.updateGUIFromParameterNode()

    def updateGUIFromParameterNode(self, caller=None, event=None):
        """
        This method is called whenever parameter node is changed.
        The module GUI is updated to show the current state of the parameter node.
        """

        if self._parameterNode is None or self._updatingGUIFromParameterNode:
            return

        # Make sure GUI changes do not call updateParameterNodeFromGUI (it could cause infinite loop)
        self._updatingGUIFromParameterNode = True
        #nastavit

        # Update node selectors and sliders
        self.ui.inputSelector.setCurrentNode(self._parameterNode.GetNodeReference("InputVolume"))
        self.ui.outputSelector.setCurrentNode(self._parameterNode.GetNodeReference("OutputVolume"))
        self.ui.NodeBoxROI.setCurrentNode(self._parameterNode.GetNodeReference("ROI"))
        self.ui.SegmentationSelector.setCurrentNode(self._parameterNode.GetNodeReference("Segmentation"))
        self.ui.filtercheckBox.checked = (self._parameterNode.GetParameter("FilterOn") == "true")
        self.ui.DownsamplingcheckBox.checked = (self._parameterNode.GetParameter("SamplingOn") == "true")


        # Update buttons states and tooltips

        self.ui.SegmentEditorWidget.setEffectNameOrder(self.activeEffets)
        self.ui.SegmentEditorWidget.unorderedEffectsVisible = False
        if self._parameterNode.GetNodeReference("InputVolume") and self._parameterNode.GetNodeReference("OutputVolume"):
            self.ui.WindowLevelButton.toolTip = "Compute output volume"
            self.ui.WindowLevelButton.enabled = True
        
            
            if self._parameterNode.GetNodeReference("ROI"):
                self.ui.CropVolumeButton.enabled = True
                self.ui.CropVolumeButton.toolTip = "Crop/downsammplig of volume"
            else:    
                self.ui.CropVolumeButton.enabled = False
                self.ui.CropVolumeButton.toolTip = "inptu, output and ROI needed"
        else:
            self.ui.WindowLevelButton.toolTip = "Select input and output volume nodes"
            self.ui.WindowLevelButton.enabled = False
            
        if self._parameterNode.GetNodeReference("InputVolume"):
            self.ui.WindowLevelButton.enabled = True
            self.ui.WindowLevelButton.toolTip = "Set window/level"
        else:
            self.ui.WindowLevelButton.toolTip = "Select input"
            self.ui.WindowLevelButton.enabled = False



        if self._parameterNode.GetParameter("FilterOn") == "true":
            self.ui.ListofFilters.enabled = True
            if self._parameterNode.GetNodeReference("InputVolume") and self._parameterNode.GetNodeReference("OutputVolume"):
                self.ui.FiltrationButton.enabled = True
            else:
                self.ui.FiltrationButton.enabled = False
                self.ui.FiltrationButton.toolTip = "Select input and output volume nodes"

        else:
            self.ui.ListofFilters.enabled = False
    

        if self._parameterNode.GetParameter("SamplingOn") == "true":
            self.ui.samplingSpinBox.enabled = True
            self.ui.groupBox_2.show()
        else:
            self.ui.samplingSpinBox.enabled = False
            self.ui.groupBox_2.hide()

        if self._parameterNode.GetNodeReference("Segmentation"):
            self.ui.SegmentsInNode.enabled = True
            self.ui.MaskSetUpButton.enabled = True
            self.ui.Mask32Button.enabled = True
            self.ui.Mask64Button.enabled = True
            self.ui.Mask128Button.enabled = True
        else:
            self.ui.SegmentsInNode.enabled = False
            self.ui.Mask32Button.enabled = False
            self.ui.Mask64Button.enabled = False
            self.ui.Mask128Button.enabled = False
            self.ui.MaskSetUpButton.enabled = False

        if self._parameterNode.GetParameter("Segment")=="None":
            self.ui.MaskSetUpButton.enabled = False
        else:
            self.ui.MaskSetUpButton.enabled = True

#not self._parameterNode.GetParameter("SizeOfMask") and

        self._updatingGUIFromParameterNode = False


    def updateParameterNodeFromGUI(self, caller=None, event=None):
        """
        This method is called when the user makes any change in the GUI.
        The changes are saved into the parameter node (so that they are restored when the scene is saved and loaded).
        """

        if self._parameterNode is None or self._updatingGUIFromParameterNode:
            return

        wasModified = self._parameterNode.StartModify()  # Modify all properties in a single batch
        #nastavit
        self._parameterNode.SetNodeReferenceID("InputVolume", self.ui.inputSelector.currentNodeID)
        self._parameterNode.SetNodeReferenceID("OutputVolume", self.ui.outputSelector.currentNodeID)
        self._parameterNode.SetNodeReferenceID("ROI", self.ui.NodeBoxROI.currentNodeID)
        self._parameterNode.SetNodeReferenceID("Segmentation",self.ui.SegmentationSelector.currentNodeID)
        self._parameterNode.SetParameter("ScaleofSampling",str(self.ui.samplingSpinBox.value))
        self._parameterNode.SetParameter("Filter", self.ui.ListofFilters.currentText)
        self._parameterNode.SetParameter("Segment", self.ui.SegmentsInNode.currentText)
        self._parameterNode.SetParameter("SamplingOn", "true" if self.ui.DownsamplingcheckBox.checked else "false")
        self._parameterNode.SetParameter("FilterOn", "true" if self.ui.filtercheckBox.checked else "false")

        self._parameterNode.EndModify(wasModified)


    #Sets parametres about which segment will be created. Parameters are set by status of checkbox
    def SegmentsForSegmentation(self):

        if self.ui.SegmentNotspecifiedCheckBox.isChecked()==True and self.ui.WholeHeartCheckBox.isChecked()==False:
            self._parameterNode.SetParameter("Unspecified","true")

            self.ui.LeftVentricleCheckBox.setChecked(False)
            self._parameterNode.SetParameter("LeftVentricle","false")

            self.ui.RightVentricleCheckBox.setChecked(False)
            self._parameterNode.SetParameter("RightVentricle","false")

            self.ui.LeftAtriumCheckBox.setChecked(False)
            self._parameterNode.SetParameter("LeftAtrium","false")

            self._parameterNode.SetParameter("RightAtrium","false")
            self.ui.RightAtriumCheckBox.setChecked(False)

        elif self.ui.WholeHeartCheckBox.isChecked()==True:
            self._parameterNode.SetParameter("LeftVentricle","true")
            self.ui.LeftVentricleCheckBox.setChecked(True)
            self._parameterNode.SetParameter("RightVentricle","true")
            self.ui.RightVentricleCheckBox.setChecked(True)
            self._parameterNode.SetParameter("LeftAtrium","true")
            self.ui.LeftAtriumCheckBox.setChecked(True)
            self._parameterNode.SetParameter("RightAtrium","true")
            self.ui.RightAtriumCheckBox.setChecked(True)
            self._parameterNode.SetParameter("Unspecified","False")
            self.ui.SegmentNotspecifiedCheckBox.setChecked(False)
        else:
            self._parameterNode.SetParameter("Unspecified","False")
            self._parameterNode.SetParameter("LeftVentricle", "true" if self.ui.LeftVentricleCheckBox.checked else "false")
            self._parameterNode.SetParameter("RightVentricle", "true" if self.ui.RightVentricleCheckBox.checked else "false")
            self._parameterNode.SetParameter("LeftAtrium", "true" if self.ui.LeftAtriumCheckBox.checked else "false")
            self._parameterNode.SetParameter("RightAtrium", "true" if self.ui.RightAtriumCheckBox.checked else "false")

    # saves data (crop image and crop segmentation in nifti format) 
    def creatTestData(self):

        NameOfsegment = []
        segment = self._parameterNode.GetParameter("Segment")
        NameOfsegment.append(segment)
        SegmentationNode = self.ui.SegmentationSelector.currentNode()

        output = slicer.mrmlScene.AddNewNodeByClass('vtkMRMLScalarVolumeNode')
        path = os.path.dirname(__file__)
        self.logic.Preprocessing(self.ui.inputSelector.currentNode(),output,self.ui.NodeBoxROI.currentNode(),1,1)
        self.logic.Creat_saveData(SegmentationNode,output,NameOfsegment,path)


        print(self._parameterNode.GetParameter("Segment"))




    # sets visibility of ROI
    def SetVisibility(self):
        roiNode = self.ui.NodeBoxROI.currentNode()
        if roiNode.GetDisplayNode().GetVisibility()== True:
            roiNode.GetDisplayNode().SetVisibility(False)
        else:
            roiNode.GetDisplayNode().SetVisibility(True)


    #function displaySelectedVolume displays selected input volume node
    #Slice Composite Node has 3 views >> red(0),green(1),yellow(2)
    #SliceViews reset automaticly adjust slice view's field of view

    # Force slicer to show selected input Volume
    def displaySelectedVolume(self):
        Volume = self.ui.inputSelector.currentNode()
        VolumeID = Volume.GetID()

        #slicer.mrmlScene.AddNewNodeByClass("vtkMRMLCropVolumeParametersNode")
        #numberOfCompositeNodes = slicer.mrmlScene.GetNumberOfNodesByClass("vtkMRMLSliceCompositeNode")
        for i in range(3):
            compositeNode = slicer.mrmlScene.GetNthNodeByClass(i,"vtkMRMLSliceCompositeNode")
            compositeNode.CreateNodeInstance()
            compositeNode.SetBackgroundVolumeID(VolumeID)
        slicer.util.resetSliceViews()

    #If segmentation is changed the list of avaible segments will be updated
    def refreshSegments(self):
        if  not self._parameterNode.GetNodeReference("Segmentation"):
            self.ui.SegmentsInNode.clear()
        else:
            self.ui.SegmentsInNode.clear()
            self.ui.SegmentsInNode.addItem("None")
            segmentationNode = self.ui.SegmentationSelector.currentNode()
            segments = segmentationNode.GetSegmentation().GetSegmentIDs()
            self.ui.SegmentsInNode.addItems(segments)
            self._parameterNode.SetParameter("Segment","None")

    # for downsamling and crop volume
    def apllyCropVolume(self):
        if self._parameterNode.GetParameter("SamplingOn") == "true":
            samplingscale = int(float(self._parameterNode.GetParameter("ScaleofSampling")))
            interpolation = int(self._parameterNode.GetParameter("UseInterpolation"))
        else:
            samplingscale = 1
            interpolation = int(self._parameterNode.GetParameter("UseInterpolation"))

        self.logic.Preprocessing(self.ui.inputSelector.currentNode(),self.ui.outputSelector.currentNode(),self.ui.NodeBoxROI.currentNode(),samplingscale,interpolation)
        print("hotovo")

    #selection of interpolation mode [NN=1, linear=2, WS=3, Bspline=4]
    def SelectInterpolation(self):
        if self.ui.radioButtonNN.isChecked()==True:
            self._parameterNode.SetParameter("UseInterpolation","1")
        if self.ui.radioButtonLinear.isChecked()==True:
            self._parameterNode.SetParameter("UseInterpolation","2")
        if self.ui.radioButtonWS.isChecked()==True:
            self._parameterNode.SetParameter("UseInterpolation","3")
        if self.ui.radioButtonBspline.isChecked()==True:
            self._parameterNode.SetParameter("UseInterpolation","4")

    # setup size of Mask
    def SelectSizeOfMask(self):
        if self.ui.Mask32Button.isChecked()==True:
            self._parameterNode.SetParameter("SizeOfMask","32")

        if self.ui.Mask64Button.isChecked()==True:
            self._parameterNode.SetParameter("SizeOfMask","64")

        if self.ui.Mask128Button.isChecked()==True:
            self._parameterNode.SetParameter("SizeOfMask","128")


    #creat Mask of size "SizeofMask" with the center in center of selected segment
    def applyCreatMaskROI(self):
        NameOfsegment = self._parameterNode.GetParameter("Segment")
        #Segment = self.ui.SegmentationSelector.currentNode().GetSegmentation().GetSegment(NameOfsegment)
        roi = self.logic.CreatMaskROI(int(self._parameterNode.GetParameter("SizeOfMask")),self.ui.SegmentationSelector.currentNode(),NameOfsegment)
        roiID = roi.GetID()
        self._parameterNode.SetNodeReferenceID("ROI", roiID)

    #aplly filtration based on selected filter
    def applyFiltration(self):
        if not self.ui.inputSelector.currentNode() or not self.ui.outputSelector.currentNode():
            message = "Input and Output needed! Please select Input a output volume."
            self.Warningbox(message)
            return
        if self._parameterNode.GetParameter("Filter") == "RankImageFilter":
            TypeOfFilter = sitk.RankImageFilter()
            TypeOfFilter.SetDebug(False)
            TypeOfFilter.SetNumberOfThreads(8)
            TypeOfFilter.SetNumberOfWorkUnits(0)
            TypeOfFilter.SetRadius((1, 1, 1))
            TypeOfFilter.SetRank(0.5)
            self.logic.BasicFiltration(self.ui.inputSelector.currentNode(),self.ui.outputSelector.currentNode(),TypeOfFilter)
        if self._parameterNode.GetParameter("Filter") == "MedianImageFilter":
            TypeOfFilter = sitk.MedianImageFilter()
            TypeOfFilter.SetRadius((1, 1, 1))
            self.logic.BasicFiltration(self.ui.inputSelector.currentNode(),self.ui.outputSelector.currentNode(),TypeOfFilter)


    # changes window/level - checkable(2 states)
    def applyWindowLevel(self):

        if self.ui.WindowLevelButton.isChecked():
        # with slicer.util.tryWithErrorDisplay("Failed to compute results.", waitCursor=True):
            self.logic.windowlevel(self.ui.inputSelector.currentNode(),900,350)
        else:
            self.ui.inputSelector.currentNode().GetDisplayNode().SetAutoWindowLevel(True)

    # creates ROI - fitting on selected input
    def applyCreatROI(self):
        if self.ui.NodeBoxROI.currentNode():
           self.ui.NodeBoxROI.currentNode().GetDisplayNode().SetVisibility(False)
        self.logic.CreatROI(self.ui.inputSelector.currentNode())

    #Changes parameter "activeSegments" - based on selected segments in GUI or when the segments is created
    def activeSegmentChangeStatus(self,Heartpart,status):
        ActiveSegments = self._parameterNode.GetParameter("activeSegments")
        if Heartpart == 1:
            ActiveSegments =ActiveSegments[:1] + status + ActiveSegments[2:]
            self._parameterNode.SetParameter("activeSegments",ActiveSegments)
        if Heartpart == 2:
            ActiveSegments = ActiveSegments[:2] + status + ActiveSegments[3:]
            self._parameterNode.SetParameter("activeSegments",ActiveSegments)
        if Heartpart == 3:
            ActiveSegments = ActiveSegments[:3] + status + ActiveSegments[4:]
            self._parameterNode.SetParameter("activeSegments",ActiveSegments)
        if Heartpart == 4:
            ActiveSegments = ActiveSegments[:4] + status
            self._parameterNode.SetParameter("activeSegments",ActiveSegments)

    # checks active segments - selected segments for creation
    def activeSegments(self):
        if self._parameterNode.GetParameter("Unspecified")=="true":
            ActiveSegments = "10000"
            self._parameterNode.SetParameter("activeSegments",ActiveSegments)
            return
        if self._parameterNode.GetParameter("LeftVentricle")=="true":
            self.activeSegmentChangeStatus(1,"1")

        if self._parameterNode.GetParameter("RightVentricle")=="true":
            self.activeSegmentChangeStatus(2,"1")

        if self._parameterNode.GetParameter("LeftAtrium")=="true":
            self.activeSegmentChangeStatus(3,"1")

        if self._parameterNode.GetParameter("RightAtrium")=="true":
            self.activeSegmentChangeStatus(4,"1")

    # Initiate segmentation
    def startSegmentation(self):
        self.ui.CollapsiblePreprocessingButton.collapsed = True
        self.ui.groupBox.hide()
        self.ui.CancleButton.enabled = True
        self.ui.StartSegmentationButton.hide()
        self.activeSegments()

        ActiveSegments = self._parameterNode.GetParameter("activeSegments")
        if ActiveSegments == "00000":
            ActiveSegments = "10000"
        int_list = [int(s) for s in ActiveSegments]
        for i in range(0,5):
            if int_list[i]==1:
                if i == 0:
                    NameOfSegment = "Unspecified"
                    ActiveSegments = "00000"
                    self._parameterNode.SetParameter("activeSegments",ActiveSegments)
                if i == 1:
                    NameOfSegment = "LeftVentricle"
                    self.activeSegmentChangeStatus(1,"0")
                if i == 2:
                    NameOfSegment = "RightVentricle"
                    self.activeSegmentChangeStatus(2,"0")
                if i == 3:
                    NameOfSegment = "LeftAtrium"
                    self.activeSegmentChangeStatus(3,"0")
                if i == 4:
                    NameOfSegment = "RightAtrium"
                    self.activeSegmentChangeStatus(4,"0")

                # calls logic function, set up segment editor for use
                CurrentSegmentation = self.logic.initiateSegmentEditor(self.ui.inputSelector.currentNode())
                self._parameterNode.SetNodeReferenceID("Segmentation",CurrentSegmentation.GetID())
                # creates first segment
                self.logic.InputForSegments(NameOfSegment,CurrentSegmentation)
                break

    # creats warning box
    def Warningbox(self,message):
        messageBox = ctk.ctkMessageBox()
        messageBox.setIcon(messageBox.Warning)
        messageBox.setText(message)
        messageBox.setStandardButtons(messageBox.Ok)
        messageBox.exec_()

    # chcecks for missing inputs
    def checkinputs(self,SegmentationNode,SourceVolumeNode):
        error = False
        if SegmentationNode == None:
            message = ("No Segmentation! Choose segmentation in SegmentEditor")
            self.Warningbox(message)
            error = True
            return
        elif SourceVolumeNode == None:
            message = ("No SourceVolume! Choose Source Volume in SegmentEditor")
            self.Warningbox(message)
            error = True
        else:
            numberOfSegments = SegmentationNode.GetSegmentation().GetNumberOfSegments()
            if numberOfSegments == 0:
                message = ("No Segments! Add segments in SegmentEditor or choose segmentation with segments.")
                self.Warningbox(message)
                error = True
                
        return error
            
            
    #crates rest of selected segments, creat background segment
    def setBackground(self):
        ActiveSegments = self._parameterNode.GetParameter("activeSegments")
        int_list = [int(s) for s in ActiveSegments]
        SegmentationNode = slicer.mrmlScene.GetNodeByID(self.ui.SegmentEditorWidget.segmentationNodeID())
        SourceVolumeNode = slicer.mrmlScene.GetNodeByID(self.ui.SegmentEditorWidget.masterVolumeNodeID())

        #check inputs
        error = self.checkinputs(SegmentationNode,SourceVolumeNode)
        if error == True:
            return
        
        #creates rest of selected segments (second segment up to 4th segment)
        if max(int_list)>0:
            for i in range(0,5):
                if int_list[i]==1:
                    if i == 1:
                        NameOfSegment = "LeftVentricle"
                        self.activeSegmentChangeStatus(1,"0")
                    if i == 2:
                        NameOfSegment = "RightVentricle"
                        self.activeSegmentChangeStatus(2,"0")
                    if i == 3:
                        NameOfSegment = "LeftAtrium"
                        self.activeSegmentChangeStatus(3,"0")
                    if i == 4:
                        NameOfSegment = "RightAtrium"
                        self.activeSegmentChangeStatus(4,"0")

                    self.logic.InputForSegments(NameOfSegment,SegmentationNode)
                    break
        else:
        # Check for inputs: no segments should be emtpy, calls function for calculating center of segments

            ListOfSegments = SegmentationNode.GetSegmentation().GetSegmentIDs()
            numberOfSegments = SegmentationNode.GetSegmentation().GetNumberOfSegments()
            CenterOfSegments = []
            for seg_id in range(0,numberOfSegments):
                SegmentNode = ListOfSegments[seg_id]
                if slicer.util.arrayFromSegment(SegmentationNode, SegmentNode).max() == 0:
                    message = ("Segment "+str(seg_id+1)+ " is empty!")
                    self.Warningbox(message)
                    return
                else:
                    Center = self.logic.GetCenter(SourceVolumeNode,SegmentationNode,SegmentNode)
                    CenterOfSegments.append(Center)
                CenterOfSegment_Ijk = self.logic.getMinMaxofCenter(CenterOfSegments)


            # Calls function which creates BG segment
            self.logic.GetBackgroundBox(SourceVolumeNode,SegmentationNode,CenterOfSegment_Ijk)

            #update status of widgets
            self.ui.StartSegmentationButton.show()
            self.ui.groupBox.show()
            self.ui.CancleButton.enabled = False
    

    # Calls function for segmentation, creates objects necessery for segmentation, deletes used segments, save parameters in case user will click update
    def Watershed(self):

        preSegmentationNode = slicer.mrmlScene.GetNodeByID(self.ui.SegmentEditorWidget.segmentationNodeID())
        SourceVolumeID = self.ui.SegmentEditorWidget.masterVolumeNodeID()
        SourceVolumeNode = slicer.mrmlScene.GetNodeByID(SourceVolumeID)
        
        error = self.checkinputs(preSegmentationNode,SourceVolumeNode)
        if error == True:
            return
        
        ListOfSegments = preSegmentationNode.GetSegmentation().GetSegmentIDs()
        IsBackground = preSegmentationNode.GetSegmentation().GetSegment("Background")

        # check if crop input volume exist (to save time when updating)
        cropinputVolume = slicer.mrmlScene.GetNodeByID(self._parameterNode.GetParameter('cropImageID'))
        if cropinputVolume == None:
            if IsBackground != None:
                preSegmentationNode.GetDisplayNode().SetSegmentVisibility("Background",True)
                roi = self.logic.BoundingBoxAndCroping(preSegmentationNode)
                roi.GetDisplayNode().SetVisibility(False)
                cropinputVolume = slicer.mrmlScene.AddNewNodeByClass('vtkMRMLScalarVolumeNode')
                slicer.modules.cropvolume.logic().CropInterpolated(roi,SourceVolumeNode,cropinputVolume,"true",1,1,0)
                slicer.mrmlScene.RemoveNode(roi)
                self.logic.thresholding(cropinputVolume,cropinputVolume,950,-50)
                cropinputVolumeID = cropinputVolume.GetID()
                self._parameterNode.SetParameter("cropImageID",cropinputVolumeID)

                cropinputVolume = self.logic.gradientfiltr(cropinputVolume,cropinputVolumeID)

            else:
                # in case no background segment
                volumesLogic = slicer.modules.volumes.logic()
                cropinputVolume = volumesLogic.CloneVolume(slicer.mrmlScene, SourceVolumeNode, "Cloned volume")
                self.logic.thresholding(cropinputVolume,cropinputVolume,950,-50)
                cropinputVolumeID = cropinputVolume.GetID()
                self._parameterNode.SetParameter("cropImageID",cropinputVolumeID)

                cropinputVolume = self.logic.gradientfiltr(cropinputVolume,cropinputVolumeID)

        else:
            # when updating, existing segmentation will be deleted
            Segmentationpreview = slicer.mrmlScene.GetNodeByID(self._parameterNode.GetParameter("SegmentationInProcessID"))
            if Segmentationpreview:
                slicer.mrmlScene.RemoveNode(Segmentationpreview)



        SegmentationNode = self.logic.WatershedSegmentation(cropinputVolume,preSegmentationNode,ListOfSegments)

         # set all segments visible, display setup
        segmentationDisplayNode = slicer.vtkMRMLSegmentationDisplayNode()
        displayNode = slicer.mrmlScene.AddNode(segmentationDisplayNode)
        SegmentationNode.SetAndObserveDisplayNodeID(displayNode.GetID())
        SegmentationNode.GetDisplayNode().SetAllSegmentsVisibility3D(True)
        SegmentationNode.GetDisplayNode().SetAllSegmentsVisibility(True)
        SegmentationNode.GetDisplayNode().SetSegmentVisibility3D("Background",False)
        SegmentationNode.GetDisplayNode().SetSegmentVisibility("Background",False)
        SegmentationNode.GetDisplayNode().Visibility3DOn()
        nodes = [slicer.mrmlScene.GetFirstNodeByName("Red").GetID(),slicer.mrmlScene.GetFirstNodeByName("View1").GetID(),slicer.mrmlScene.GetFirstNodeByName("Green").GetID(),slicer.mrmlScene.GetFirstNodeByName("Yellow").GetID()]
        SegmentationNode.GetDisplayNode().SetViewNodeIDs(nodes)


        self.ui.SegmentEditorWidget.setSegmentationNode(preSegmentationNode)
        self.ui.SegmentEditorWidget.setSourceVolumeNode(SourceVolumeNode)

        # show buttons for update and end segmentation
        if self._parameterNode.GetParameter("updateHidden")== "true":
            self.ui.finalSegmentationgroupBox.show()
            self.ui.SegmentationgroupBox.hide()
            self._parameterNode.SetParameter("updateHidden","false")
        SegmentationID = SegmentationNode.GetID()

        if self._parameterNode.GetParameter("SegmentationInProcessID") != "None":
            SegmentationForDel = slicer.mrmlScene.GetNodeByID(self._parameterNode.GetParameter("SegmentationInProcessID"))
            slicer.mrmlScene.RemoveNode(SegmentationForDel)

        # set segments to prewiev ...user dont need to change in segment editor and edit input immedietly
        self._parameterNode.SetParameter("SegmentationInProcessID",SegmentationID)
        self._parameterNode.SetParameter("SourceVolumeID",SourceVolumeID)
        

    # post-processing and clean up
    def endSegmentation(self):
        FinalSegmentation = slicer.mrmlScene.GetNodeByID(self._parameterNode.GetParameter("SegmentationInProcessID"))
        SourceVolume = slicer.mrmlScene.GetNodeByID(self._parameterNode.GetParameter("SourceVolumeID"))

        self.logic.DilatationErosion(SourceVolume,FinalSegmentation)

        CropVolume = slicer.mrmlScene.GetNodeByID(self._parameterNode.GetParameter('cropImageID'))
        slicer.mrmlScene.RemoveNode(CropVolume)
        self._parameterNode.SetParameter("cropImageID","None")

        self.ui.SegmentEditorWidget.setSegmentationNode(FinalSegmentation)
        self.ui.SegmentEditorWidget.setSourceVolumeNode(SourceVolume)
        self._parameterNode.SetParameter("SegmentationInProcessID","None")

        self._parameterNode.SetParameter("updateHidden","true")
        self.ui.SegmentationgroupBox.show()
        self.ui.finalSegmentationgroupBox.hide()


    #ends contructor during process
    def Cancel_stopInitalize(self):

        segmentationNode = self._parameterNode.GetNodeReference("Segmentation")
        slicer.mrmlScene.RemoveNode(segmentationNode)
        self.ui.StartSegmentationButton.show()
        self.ui.CancleButton.enabled = False
        self.ui.groupBox.show()
        self._parameterNode.SetParameter("updateHidden","true")


#
# segment1Logic
#

class HeartSegmentationLogic(ScriptedLoadableModuleLogic):
    """This class should implement all the actual
    computation done by your module.  The interface
    should be such that other python code can import
    this class and make use of the functionality without
    requiring an instance of the Widget.
    Uses ScriptedLoadableModuleLogic base class, available at:
    https://github.com/Slicer/Slicer/blob/main/Base/Python/slicer/ScriptedLoadableModule.py
    """

    def __init__(self):
        """
        Called when the logic class is instantiated. Can be used for initializing member variables.
        """
        ScriptedLoadableModuleLogic.__init__(self)

        self.outputSegmentation = None
        self.segmentEditorWidget = None

    def setDefaultParameters(self, parameterNode):

        """
        Initialize parameter node with default settings.
        """
        if not parameterNode.GetParameter("Window"):
            parameterNode.SetParameter("Window", "400")
            if not parameterNode.GetParameter("Level"):
                parameterNode.SetParameter("Level", "350")

        if not parameterNode.GetParameter("SizeOfMask"):
            parameterNode.SetParameter("SizeOfMask", "32")

        if not parameterNode.GetParameter("FilterOn"):
            parameterNode.SetParameter("FilterOn", "false")

        if not parameterNode.GetParameter("ScaleOfSampling"):
            parameterNode.SetParameter("ScaleOfSampling", "1.0")

        if not parameterNode.GetParameter("Segment"):
            parameterNode.SetParameter("Segment", "None")

        if not parameterNode.GetParameter("Segmentation"):
            parameterNode.SetParameter("Segmentation", "None")

    
    def segmentationNode(self):
        return self.getParameterNode().GetNodeReference("Segmentation")
    
    # Gets min and max of each segment and then takes max and min in each axis
    def getMinMaxofCenter(self,CenterOfSegments):
        max_x = -float('inf')
        max_y = -float('inf')
        max_z = -float('inf')
        min_x = float('inf')
        min_y = float('inf')
        min_z = float('inf')
        for center in CenterOfSegments:
            max_x = max(max_x, center[0])
            max_y = max(max_y, center[1])
            max_z = max(max_z, center[2])
            min_x = min(min_x, center[0])
            min_y = min(min_y, center[1])
            min_z = min(min_z, center[2])
    
        return [(max_x,max_y,max_z),(min_x,min_y,min_z)]

    # creats bounding box around Segment
    def BoundingBoxAndCroping(self, segmentationNode):
        import SegmentStatistics
        segStatLogic = SegmentStatistics.SegmentStatisticsLogic()
        segStatLogic.getParameterNode().SetParameter("Segmentation", segmentationNode.GetID())
        segStatLogic.getParameterNode().SetParameter("LabelmapSegmentStatisticsPlugin.obb_origin_ras.enabled",str(True))
        segStatLogic.getParameterNode().SetParameter("LabelmapSegmentStatisticsPlugin.obb_diameter_mm.enabled",str(True))
        segStatLogic.getParameterNode().SetParameter("LabelmapSegmentStatisticsPlugin.obb_direction_ras_x.enabled",str(True))
        segStatLogic.getParameterNode().SetParameter("LabelmapSegmentStatisticsPlugin.obb_direction_ras_y.enabled",str(True))
        segStatLogic.getParameterNode().SetParameter("LabelmapSegmentStatisticsPlugin.obb_direction_ras_z.enabled",str(True))
        segStatLogic.computeStatistics()
        stats = segStatLogic.getStatistics()

        segmentId = "Background"
        # segmentId = slicer.mrmlScene.GetNodeByID(segment)
        obb_origin_ras = np.array(stats[segmentId,"LabelmapSegmentStatisticsPlugin.obb_origin_ras"])
        obb_diameter_mm = np.array(stats[segmentId,"LabelmapSegmentStatisticsPlugin.obb_diameter_mm"])
        obb_direction_ras_x = np.array(stats[segmentId,"LabelmapSegmentStatisticsPlugin.obb_direction_ras_x"])
        obb_direction_ras_y = np.array(stats[segmentId,"LabelmapSegmentStatisticsPlugin.obb_direction_ras_y"])
        obb_direction_ras_z = np.array(stats[segmentId,"LabelmapSegmentStatisticsPlugin.obb_direction_ras_z"])

        segment = segmentationNode.GetSegmentation().GetSegment(segmentId)
        roi=slicer.mrmlScene.AddNewNodeByClass("vtkMRMLMarkupsROINode")
        roi.SetName(segment.GetName())
        roi.SetSize(obb_diameter_mm)

        obb_center_ras = obb_origin_ras+0.5*(obb_diameter_mm[0] * obb_direction_ras_x + obb_diameter_mm[1] * obb_direction_ras_y + obb_diameter_mm[2] * obb_direction_ras_z)
        boundingBoxToRasTransform = np.row_stack((np.column_stack((obb_direction_ras_x, obb_direction_ras_y, obb_direction_ras_z, obb_center_ras)), (0, 0, 0, 1)))
        boundingBoxToRasTransformMatrix = slicer.util.vtkMatrixFromArray(boundingBoxToRasTransform)
        roi.SetAndObserveObjectToNodeMatrix(boundingBoxToRasTransformMatrix)
        return roi

    #Creats ROI around specific segment with size SizeOfMask
    def CreatMaskROI(self,SizeOfMask,SegmentNode,Segment):
        center = SegmentNode.GetSegmentCenter(Segment)
        roiNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLMarkupsROINode")
        cropVolumeParameters = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLCropVolumeParametersNode")
        roiNode.SetSize(SizeOfMask,SizeOfMask,SizeOfMask)
        roiNode.SetCenter(center)
        #roiNode.GetDisplayNode().SetHandlesInteractive(False)
        slicer.modules.cropvolume.logic().SnapROIToVoxelGrid(cropVolumeParameters)
        return roiNode

    # creats roi, ROI is size of Volume
    def CreatROI(self, inputVolume):
        if not inputVolume:
            raise ValueError("input needed")

        roiNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLMarkupsROINode")
        cropVolumeParameters = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLCropVolumeParametersNode")
        cropVolumeParameters.SetInputVolumeNodeID(inputVolume.GetID())
        cropVolumeParameters.SetROINodeID(roiNode.GetID())
        slicer.modules.cropvolume.logic().SnapROIToVoxelGrid(cropVolumeParameters)  # optional (rotates the ROI to match the volume axis directions)
        slicer.modules.cropvolume.logic().FitROIToInputVolume(cropVolumeParameters)


    def WatershedSegmentation(self,inputVolume,SegmentNode,Segments):
        qt.QApplication.setOverrideCursor(qt.Qt.WaitCursor)

        volumenode = sitkUtils.PullVolumeFromSlicer(inputVolume)

        # ObjectScalemm = 2.3

        # #creats gradient image and aplies watershed
        # gradientMagnitude = sitk.GradientMagnitudeRecursiveGaussian(volumenode,ObjectScalemm)
        labelmapVolumeNode = slicer.mrmlScene.AddNewNodeByClass('vtkMRMLLabelMapVolumeNode')
        slicer.modules.segmentations.logic().ExportSegmentsToLabelmapNode(SegmentNode,Segments,labelmapVolumeNode,inputVolume)
        labelImage = sitkUtils.PullVolumeFromSlicer(labelmapVolumeNode)
        waterShedSegmentation = sitk.MorphologicalWatershedFromMarkers(inputVolume, labelImage, markWatershedLine=False, fullyConnected=False)
        sitkUtils.PushVolumeToSlicer(waterShedSegmentation,labelmapVolumeNode)
        HeartSegmentation = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentationNode","Heart_Segmentation")

        slicer.modules.segmentations.logic().ImportLabelmapToSegmentationNode(labelmapVolumeNode,HeartSegmentation,"")
        slicer.mrmlScene.RemoveNode(labelmapVolumeNode)

        qt.QApplication.restoreOverrideCursor()
        return HeartSegmentation


    def GetBackgroundBox(self,inputVolume,SegmentNode,CenterOfSegment_Ijk):
        #end segmenteditor effect
        if self.segmentEditorWidget:
            self.segmentEditorWidget.setActiveEffect(None)

        # get centers
        center_max = CenterOfSegment_Ijk[0]
        center_min = CenterOfSegment_Ijk[1]

        max_x =int(round(center_max[0]))
        max_y =int(round(center_max[1]))
        min_x = int(round(center_min[0]))
        min_y =int(round(center_min[1]))
        max_z =int(round(center_max[2]))
        min_z = int(round(center_min[2]))
        diff_z = int(round(max_z)) - int(round(min_z))
        diff_x = int(round(max_x)) - int(round(min_x))
        z = int(round(min_z) + round(diff_z/2))
        x = int(round(min_x) + round(diff_x/2))
        spacing = inputVolume.GetSpacing()

        x_AddValue = round(55/spacing[0])
        y_AddValue = round(55/spacing[1])
        z_AddValue = round(55/spacing[1])


        Z = range(min_z-z_AddValue,max_z+z_AddValue,8)
        X = range(min_x  -x_AddValue,max_x + x_AddValue,12)
        Y = range(min_y - y_AddValue,max_y + y_AddValue,12)
        append = vtk.vtkAppendPolyData()

        # creats background with vtkSphereSource
        for n in X:
            backgroundSeed = vtk.vtkSphereSource()
            backgroundSeed.SetCenter([n,max_y + y_AddValue,z])
            backgroundSeed.SetRadius(10)
            backgroundSeed.Update()
            append.AddInputData(backgroundSeed.GetOutput())
            backgroundSeed = vtk.vtkSphereSource()
            backgroundSeed.SetCenter([n,min_y - y_AddValue,z])
            backgroundSeed.SetRadius(10)
            backgroundSeed.Update()
            append.AddInputData(backgroundSeed.GetOutput())

        for i in Y:
            backgroundSeed = vtk.vtkSphereSource()
            backgroundSeed.SetCenter([max_x + x_AddValue,i,z])
            backgroundSeed.SetRadius(10)
            backgroundSeed.Update()
            append.AddInputData(backgroundSeed.GetOutput())
            backgroundSeed = vtk.vtkSphereSource()
            backgroundSeed.SetCenter([min_x - x_AddValue,i,z])
            backgroundSeed.SetRadius(10)
            backgroundSeed.Update()
            append.AddInputData(backgroundSeed.GetOutput())

            backgroundSeed = vtk.vtkSphereSource()
            backgroundSeed.SetCenter([x,i, max_z+ z_AddValue])
            backgroundSeed.SetRadius(10)
            backgroundSeed.Update()
            append.AddInputData(backgroundSeed.GetOutput())
            backgroundSeed = vtk.vtkSphereSource()
            backgroundSeed.SetCenter([x,i, min_z - z_AddValue])
            backgroundSeed.SetRadius(10)
            backgroundSeed.Update()
            append.AddInputData(backgroundSeed.GetOutput())

        for k in Z:
            backgroundSeed = vtk.vtkSphereSource()
            backgroundSeed.SetCenter([x,max_y + y_AddValue,k])
            backgroundSeed.SetRadius(10)
            backgroundSeed.Update()
            append.AddInputData(backgroundSeed.GetOutput())
            backgroundSeed = vtk.vtkSphereSource()
            backgroundSeed.SetCenter([x,min_y - y_AddValue,k])
            backgroundSeed.SetRadius(10)
            backgroundSeed.Update()
            append.AddInputData(backgroundSeed.GetOutput())


        #tranformation back to RAS
        append.Update()
        conversionNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentationNode", "ConversionNode")
        conversionNode.AddSegmentFromClosedSurfaceRepresentation(append.GetOutput(), "Background", [0.0,1.0,0.0])
        ijkToRas = vtk.vtkMatrix4x4()
        inputVolume.GetIJKToRASMatrix(ijkToRas)
        conversionNode.ApplyTransformMatrix(ijkToRas)
        convertedSegment = conversionNode.GetSegmentation().GetSegment("Background")
        SegmentNode.GetSegmentation().AddSegment(convertedSegment)
        slicer.mrmlScene.RemoveNode(conversionNode)

        #Gets centers of Segments, tranforms RAS to IJK
    def GetCenter(self,inputVolume,SegmentNode,Segment):
        RasToIjk = vtk.vtkMatrix4x4()
        inputVolume.GetRASToIJKMatrix(RasToIjk)
        segmentCenter_RAS = SegmentNode.GetSegmentCenterRAS(Segment)
        center = np.array(segmentCenter_RAS)
        center = np.append(center, 1.0)
        point_Ijk = [0, 0, 0, 1]
        RasToIjk.MultiplyPoint(center, point_Ijk)
        Center_Ijk = point_Ijk[:3]
        return Center_Ijk

    # apply filtration, input must be itk image
    def BasicFiltration(self,inputVolume,outputVolume,TypeOfFilter):
        inputimage = sitkUtils.PullVolumeFromSlicer(inputVolume)
        outputImage = TypeOfFilter.Execute(inputimage)
        sitkUtils.PushVolumeToSlicer(outputImage, outputVolume)
        VolumeID = outputVolume.GetID()
        numberOfCompositeNodes = slicer.mrmlScene.GetNumberOfNodesByClass("vtkMRMLSliceCompositeNode")
        for i in range(numberOfCompositeNodes):
            compositeNode = slicer.mrmlScene.GetNthNodeByClass(i,"vtkMRMLSliceCompositeNode")
            compositeNode.SetBackgroundVolumeID(VolumeID)
        print("Filtration done")

    #Uses cropvolume logic
    def Preprocessing(self,inputVolume,outputVolume,roiNode,samplingscale,interpolation):

        slicer.modules.cropvolume.logic().CropInterpolated(roiNode,inputVolume,outputVolume,"true",samplingscale,interpolation,0)
        #input(6)=1-4 (int) type of interpolation

    #sets window/level
    def windowlevel(self,inputVolume ,window, level, showResult=True):

        inputVolume.GetDisplayNode().SetAutoWindowLevel(False)
        inputVolume.GetDisplayNode().SetWindowLevel(window,level)

    # creats segment editor in logic class, setup segment editor
    def initiateSegmentEditor(self,inputVolume):
        self.segmentationNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentationNode", "preHeart Segmentation")
        if self.segmentEditorWidget == None:
            self.segmentEditorWidget = slicer.qMRMLSegmentEditorWidget()
            self.segmentEditorWidget.setMRMLScene(slicer.mrmlScene)
            self.segmentEditorNode = slicer.mrmlScene.GetSingletonNode("SegmentEditor", "vtkMRMLSegmentEditorNode")
            self.segmentEditorWidget.setMRMLSegmentEditorNode(self.segmentEditorNode)
        
        self.segmentEditorWidget.setSegmentationNode(self.segmentationNode)
        self.segmentEditorWidget.setSourceVolumeNode(inputVolume)

        self.segmentationNode.CreateDefaultDisplayNodes()

        CurrentSegmentation = self.segmentationNode
        return CurrentSegmentation
    
    #runs thresholding
    def thresholding(self, inputVolume, outputVolume, upper,lower):
        cliParams = {
            'InputVolume': inputVolume.GetID(),
            'OutputVolume': outputVolume.GetID(),
            'Upper': upper,
            'Lower':lower,
            'ThresholdType': 'Outside'
        }
        cliNode = slicer.cli.run(slicer.modules.thresholdscalarvolume, None, cliParams, wait_for_completion=True,update_display=False)
        slicer.mrmlScene.RemoveNode(cliNode)
    
    #Post-processing, applying Margin and Islind effects
    def DilatationErosion(self,inputVolume,SegmentationNode):

        if self.segmentEditorWidget == None:
            self.segmentEditorWidget = slicer.qMRMLSegmentEditorWidget()
            self.segmentEditorWidget.setMRMLScene(slicer.mrmlScene)
            self.segmentEditorNode = slicer.mrmlScene.GetSingletonNode("SegmentEditor", "vtkMRMLSegmentEditorNode")
            self.segmentEditorWidget.setMRMLSegmentEditorNode(self.segmentEditorNode)
            self.segmentEditorWidget.setSegmentationNode(SegmentationNode)
            self.segmentEditorWidget.setSourceVolumeNode(inputVolume) 
        else:
            self.segmentEditorWidget.setSegmentationNode(SegmentationNode)
            self.segmentEditorWidget.setSourceVolumeNode(inputVolume) 

        self.segmentEditorNode.SetOverwriteMode(slicer.vtkMRMLSegmentEditorNode.OverwriteNone)
      
        self.segmentEditorWidget.setActiveEffectByName("Margin")
        effect = self.segmentEditorWidget.activeEffect()
        effect.setParameter("ApplyToAllVisibleSegments",1)
        effect.setParameter("MarginSizeMm",str(4.00))
        effect.self().onApply()
        effect.setParameter("MarginSizeMm",str(-4.00))
        effect.self().onApply()

        self.segmentEditorWidget.setActiveEffectByName("Islands")
        effect = self.segmentEditorWidget.activeEffect()
        effect.setParameter("Operation", "KEEP_LARGEST_ISLAND")
        effect.self().onApply()

        self.segmentEditorWidget.setActiveEffectByName("Smoothing")
        effect = self.segmentEditorWidget.activeEffect()
        effect.self().onApply()

        self.segmentEditorWidget.setActiveEffect(None)

    #creats segments in segmentation
    def InputForSegments(self,NameOfSegment,SegmentationNode):

        segmentID = SegmentationNode.GetSegmentation().AddEmptySegment(NameOfSegment)

        self.segmentEditorWidget.setActiveEffectByName("Paint")
        self.segmentEditorWidget.setCurrentSegmentID(segmentID)


    def Creat_saveData(self,SegmentationNode,outputNode,NameOfsegment,path):
        
        properties = {'useCompression': 1} # use compression
        file_path = os.path.join(os.path.dirname(__file__),"ImageData.nii.gz")
        slicer.util.saveNode(outputNode, file_path, properties)

        extension = "nii.gz"
        useCompression = "true"

        segmentIds = vtk.vtkStringArray()
        for segmentName in NameOfsegment:
            segmentId = SegmentationNode.GetSegmentation().GetSegmentIdBySegmentName(segmentName)
            segmentIds = vtk.vtkStringArray()
            segmentIds.InsertNextValue(segmentId)

        slicer.vtkSlicerSegmentationsModuleLogic.ExportSegmentsBinaryLabelmapRepresentationToFiles(path, SegmentationNode, segmentIds, extension, useCompression)
        # slicer.mrmlScene.RemoveNode(labelmapVolumeNode)
        slicer.mrmlScene.RemoveNode(outputNode)

    def gradientfiltr(self, cropinput,cropinputVolumeID):
        volumenode = sitkUtils.PullVolumeFromSlicer(cropinput)
        ObjectScalemm = 2.3
                #creats gradient image and aplies watershed
        gradientMagnitude = sitk.GradientMagnitudeRecursiveGaussian(volumenode,ObjectScalemm)
        sitkUtils.PushVolumeToSlicer(gradientMagnitude,cropinputVolumeID)





class HeartSegmentationTest(ScriptedLoadableModuleTest):
    """
    This is the test case for your scripted module.
    Uses ScriptedLoadableModuleTest base class, available at:
    https://github.com/Slicer/Slicer/blob/main/Base/Python/slicer/ScriptedLoadableModule.py
    """

    def setUp(self):
        """ Do whatever is needed to reset the state - typically a scene clear will be enough.
        """
        slicer.mrmlScene.Clear()

    def runTest(self):
        """Run as few or as many tests as needed here.
        """
        self.setUp()
        self.test_HeartSegmentation1()

    def test_HeartSegmentation1(self):
        """ Ideally you should have several levels of tests.  At the lowest level
        tests should exercise the functionality of the logic with different inputs
        (both valid and invalid).  At higher levels your tests should emulate the
        way the user would interact with your code and confirm that it still works
        the way you intended.
        One of the most important features of the tests is that it should alert other
        developers when their changes will have an impact on the behavior of your
        module.  For example, if a developer removes a feature that you depend on,
        your test should break so they know that the feature is needed.
        """

        self.delayDisplay("Starting the test")

        # Get/create input data

        import SampleData
        registerSampleData()
        inputVolume = SampleData.SampleDataLogic().downloadMRHead()
        self.delayDisplay('Loaded test data set')

        outputVolume = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLScalarVolumeNode")

        # Test the module logic

        logic = HeartSegmentationLogic()

        logic.initiateSegmentEditor(inputVolume)

        SegmentationNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentationNode")

        append = vtk.vtkAppendPolyData()
        backgroundSeed = vtk.vtkSphereSource()
        backgroundSeed.SetCenter([50,50 + 50,50])
        backgroundSeed.SetRadius(10)
        backgroundSeed.Update()
        append.AddInputData(backgroundSeed.GetOutput())
        append.Update()

        CenterOfSegments = []
    

        SegmentationNode.AddSegmentFromClosedSurfaceRepresentation(append.GetOutput(), "segment", [1.0,0.0,0.0])

        center = logic.GetCenter(inputVolume,SegmentationNode,"segment")
        CenterOfSegments.append(center)

        c = logic.getMinMaxofCenter(CenterOfSegments)

        logic.GetBackgroundBox(inputVolume,SegmentationNode,c)



        self.delayDisplay('Test passed')
