import logging
import os

import qt
import vtk

import slicer

from SegmentEditorEffects import *


class SegmentEditorEffect(AbstractScriptedSegmentEditorEffect):
    """This effect uses Watershed algorithm to partition the input volume"""

    def __init__(self, scriptedEffect):
        scriptedEffect.name = "ScriptedSegmentEditorEffectModuleTemplate"
        scriptedEffect.perSegment = False  # this effect operates on all segments at once (not on a single selected segment)
        scriptedEffect.requireSegments = True  # this effect requires segment(s) existing in the segmentation
        AbstractScriptedSegmentEditorEffect.__init__(self, scriptedEffect)

    def clone(self):
        # It should not be necessary to modify this method
        import qSlicerSegmentationsEditorEffectsPythonQt as effects

        clonedEffect = effects.qSlicerSegmentEditorScriptedEffect(None)
        clonedEffect.setPythonSource(__file__.replace("\\", "/"))
        return clonedEffect

    def icon(self):
        # It should not be necessary to modify this method
        iconPath = os.path.join(os.path.dirname(__file__), "SegmentEditorEffect.png")
        if os.path.exists(iconPath):
            return qt.QIcon(iconPath)
        return qt.QIcon()

    def helpText(self):
        return """Existing segments are grown to fill the image.
The effect is different from the Grow from seeds effect in that smoothness of structures can be defined, which can prevent leakage.
To segment a single object, create a segment and paint inside and create another segment and paint outside on each axis.
"""

    def setupOptionsFrame(self):
        # Object scale slider
        self.objectScaleMmSlider = slicer.qMRMLSliderWidget()
        self.objectScaleMmSlider.setMRMLScene(slicer.mrmlScene)
        self.objectScaleMmSlider.quantity = "length"  # get unit, precision, etc. from MRML unit node
        self.objectScaleMmSlider.minimum = 0
        self.objectScaleMmSlider.maximum = 10
        self.objectScaleMmSlider.value = 2.0
        self.objectScaleMmSlider.setToolTip("Increasing this value smooths the segmentation and reduces leaks. This is the sigma used for edge detection.")
        self.scriptedEffect.addLabeledOptionsWidget("Object scale:", self.objectScaleMmSlider)
        self.objectScaleMmSlider.connect("valueChanged(double)", self.updateMRMLFromGUI)

        # Apply button
        self.applyButton = qt.QPushButton("Apply")
        self.applyButton.objectName = self.__class__.__name__ + "Apply"
        self.applyButton.setToolTip("Accept previewed result")
        self.scriptedEffect.addOptionsWidget(self.applyButton)
        self.applyButton.connect("clicked()", self.onApply)

    def createCursor(self, widget):
        # Turn off effect-specific cursor for this effect
        return slicer.util.mainWindow().cursor

    def setMRMLDefaults(self):
        self.scriptedEffect.setParameterDefault("ObjectScaleMm", 2.0)

    def updateGUIFromMRML(self):
        objectScaleMm = self.scriptedEffect.doubleParameter("ObjectScaleMm")
        wasBlocked = self.objectScaleMmSlider.blockSignals(True)
        self.objectScaleMmSlider.value = abs(objectScaleMm)
        self.objectScaleMmSlider.blockSignals(wasBlocked)

    def updateMRMLFromGUI(self):
        self.scriptedEffect.setParameter("ObjectScaleMm", self.objectScaleMmSlider.value)

    def onApply(self):
        # Get list of visible segment IDs, as the effect ignores hidden segments.
        segmentationNode = self.scriptedEffect.parameterSetNode().GetSegmentationNode()
        visibleSegmentIds = vtk.vtkStringArray()
        segmentationNode.GetDisplayNode().GetVisibleSegmentIDs(visibleSegmentIds)
        if visibleSegmentIds.GetNumberOfValues() == 0:
            logging.info("Smoothing operation skipped: there are no visible segments")
            return

        # This can be a long operation - indicate it to the user
        qt.QApplication.setOverrideCursor(qt.Qt.WaitCursor)

        # Allow users revert to this state by clicking Undo
        self.scriptedEffect.saveStateForUndo()

        # Export source image data to temporary new volume node.
        # Note: Although the original source volume node is already in the scene, we do not use it here,
        # because the source volume may have been resampled to match segmentation geometry.
        sourceVolumeNode = slicer.vtkMRMLScalarVolumeNode()
        slicer.mrmlScene.AddNode(sourceVolumeNode)
        sourceVolumeNode.SetAndObserveTransformNodeID(segmentationNode.GetTransformNodeID())
        slicer.vtkSlicerSegmentationsModuleLogic.CopyOrientedImageDataToVolumeNode(self.scriptedEffect.sourceVolumeImageData(), sourceVolumeNode)
        # Generate merged labelmap of all visible segments, as the filter expects a single labelmap with all the labels.
        mergedLabelmapNode = slicer.vtkMRMLLabelMapVolumeNode()
        slicer.mrmlScene.AddNode(mergedLabelmapNode)
        slicer.vtkSlicerSegmentationsModuleLogic.ExportSegmentsToLabelmapNode(segmentationNode, visibleSegmentIds, mergedLabelmapNode, sourceVolumeNode)

        # Run segmentation algorithm
        import SimpleITK as sitk
        import sitkUtils

        # Read input data from Slicer into SimpleITK
        labelImage = sitkUtils.PullVolumeFromSlicer(mergedLabelmapNode)
        backgroundImage = sitkUtils.PullVolumeFromSlicer(sourceVolumeNode)
        # Run watershed filter
        featureImage = sitk.GradientMagnitudeRecursiveGaussian(backgroundImage, float(self.scriptedEffect.doubleParameter("ObjectScaleMm")))
        del backgroundImage
        f = sitk.MorphologicalWatershedFromMarkersImageFilter()
        f.SetMarkWatershedLine(False)
        f.SetFullyConnected(False)
        labelImage = f.Execute(featureImage, labelImage)
        del featureImage
        # Pixel type of watershed output is the same as the input. Convert it to int16 now.
        if labelImage.GetPixelID() != sitk.sitkInt16:
            labelImage = sitk.Cast(labelImage, sitk.sitkInt16)
        # Write result from SimpleITK to Slicer
        sitkUtils.PushVolumeToSlicer(labelImage, mergedLabelmapNode)
        mergedLabelmapNode.GetImageData().Modified()
        mergedLabelmapNode.Modified()

        # Update segmentation from labelmap node and remove temporary nodes
        slicer.vtkSlicerSegmentationsModuleLogic.ImportLabelmapToSegmentationNode(mergedLabelmapNode, segmentationNode, visibleSegmentIds)
        slicer.mrmlScene.RemoveNode(sourceVolumeNode)
        slicer.mrmlScene.RemoveNode(mergedLabelmapNode)

        qt.QApplication.restoreOverrideCursor()
