#==============================================================================
#
# Project:   SharpImage
# Module:    FastMarching.py
# Language:  IronPython
# Author:    Dan Mueller
# Date:      $Date: 2007-07-06 10:57:00 +1000 (Fri, 06 Jul 2007) $
# Revision:  $Revision: 2 $
#
# Copyright (c) Queensland University of Technology (QUT) 2007.
# All rights reserved.
#
# This software is distributed WITHOUT ANY WARRANTY; without even
# the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
# PURPOSE.  See the above copyright notices for more information.
#
#==============================================================================

# Import the base script class
import ImageToImageScript
from ImageToImageScript import *

# Add reference and import required libraries
clr.AddReference("ManagedITK.LevelSetFilters")
clr.AddReference("ManagedITK.ThresholdFilters")
from itk import *

class FastMarchingScript(ImageToImageScriptObject):
    # -------------------------------------------------------------------------
    Name = "FastMarching"
    Help = """Fast marching solves an Eikonal equation where the speed is always
              non-negative and depends on the position only. Starting from an
              initial position on the front, fast marching systematically moves
              the front forward one grid point at a time. If the Speed is not
              provided, the SpeedConstant is used instead. A form is displayed to
              specify starting seed positions if the initial image is None."""
    Parameters = """(bool) Speed = True and the input is the speed image, else the speed constant is used. (True)
                    (string) Initial = the search pattern of the initial contour. (None)
                    (double) StoppingValue = Terminates when smallest trial point greater than the stopping value. (500.0)
                    (double) SpeedConstant = The speed value for the algorithm, if input is not a speed image. (1.0)"""
    Speed = True
    Initial = None
    StoppingValue = 500.0
    SpeedConstant = 1.0
    Form = None
    Seeds = None
    # -------------------------------------------------------------------------

    def Run(self):
        """ The entry-point for this script. """
        self.Initialise()

    def Initialise(self):
        """ Initialise the environment for running this script. """
        """ NOTE: This function is invoked on the main UI thread. """
        # Init base
        ImageToImageScriptObject.Initialise( self )

        # Get the initial image
        if (self.Initial != None):
            self.Initial = self.GetOpenImageMatchingPattern( self.Initial )
            if (self.Initial.PixelType.TypeAsEnum != itkPixelTypeEnum.UnsignedChar):
                raise ApplicationException( "Initial must have unsigned char pixel type." )

        # Prompt the user for seeds (if no initial image)
        if (self.Initial == None):
            self.Form = siFormSelectSeeds( self.Renderer )
            self.Form.Continue += self.Continue
            self.Form.Cancel += self.Cancel
            self.ParentApplication.AddTool( self.Form )

        # Start processing
        self.StartedWork()
        self.ParentApplication.SetApplicationAsReady( 0 )
        if (self.Initial != None): self.DoWork()
        else: self.WriteLineToConsole( "Select trial points..." )

    def Continue(self, sender, args):
        """ Allow processing to continue. """
        self.Seeds = self.Form.PointsSelection.GetIndicies( self.Input )
        self.DoWork()

    def Cancel(self, sender, args):
        """ Cancel the script. """
        self.WriteLineToConsole( "Script cancelled by user..." )
        self.UnregisterEventHandlers()
        self.FinishedWork( True )

    def ThreadedDoWork(self):
        """ Perform the main functions of the script. """
        """ NOTE: This function is invoked on a background thread. """
        try:
            # Setup
            self.ParentApplication.SetApplicationAsWorking()
            self.WriteInputName()
            
            # Create filter
            filterFastMarching = itkFastMarchingImageFilter.New( self.Output, self.Input )
            self.AddEventHandlersToProcessObject( filterFastMarching )
            filterFastMarching.StoppingValue = self.StoppingValue

            # Write parameters to console
            self.WriteLineToConsole( "StoppingValue=" + self.StoppingValue.ToString() )
            self.WriteLineToConsole( "Speed=" + self.Speed.ToString() )

            # Set speed as image or constant
            if (self.Speed):
                filterFastMarching.SetInput( self.Input )
            else:
                filterFastMarching.SpeedConstant = self.SpeedConstant
                filterFastMarching.OverrideOutputInformation = True
                filterFastMarching.OutputSize = self.Input.Size
                filterFastMarching.OutputSpacing = self.Input.Spacing
                filterFastMarching.OutputOrigin = self.Input.Origin
                self.WriteLineToConsole( "SpeedConstant=" + self.SpeedConstant.ToString() )

            # Add trial/alive points
            self.ParentApplication.SetApplicationStatusLabel( "Adding initial points..." )
            if (self.Initial == None):
                # Set trial points from seeds
                trialPoints = List[itkLevelSetNode]()
                for i in range(self.Seeds.Count):
                    trialPoints.Add( itkLevelSetNode(0.0, self.Seeds[i]) )
                    self.WriteLineToConsole( "Added trial point=" + self.Seeds[i].ToString() )
                filterFastMarching.TrialPoints = trialPoints.ToArray()
            else:
                # Set trial points from image
                self.WriteImageName( "Initial=", self.Initial )
                filterFastMarching.SetLabelImageFromContour( self.Initial )

            # Threshold FastMarching output
            filterThreshold = itkThresholdImageFilter.New( self.Output )
            self.AddEventHandlersToProcessObject( filterThreshold )
            filterThreshold.SetInput( filterFastMarching.GetOutput() )
            filterThreshold.OutsideValue = itkPixel( self.Output.PixelType, self.StoppingValue )
            pixelLower = itkPixel( self.Output.PixelType, 0.0 )
            pixelUpper = itkPixel( self.Output.PixelType, self.StoppingValue )
            filterThreshold.ThresholdOutside( pixelLower, pixelUpper )
            filterThreshold.GetOutput( self.Output )
            filterThreshold.UpdateLargestPossibleRegion()

            # Finish up
            self.DisconnectInputAndOutput()
            self.DisposeOfObject( filterFastMarching )
            self.DisposeOfObject( filterThreshold )
            self.WriteOutputName()
            self.FinishedWork( True )

        except Exception, ex:
            self.HandleException( ex )
            self.FinishedWork( False )
            self.Finalise( )

    def Finalise(self):
        """ Finalise the environment after running this script. """
        """ NOTE: This function is invoked on the main UI thread. """
        self.UnregisterEventHandlers()
        ImageToImageScriptObject.Finalise(self)

    def UnregisterEventHandlers(self):
        """ Unregister all the event handlers. """
        if (self.Form != None):
            self.Form.Continue -= self.Continue
            self.Form.Cancel -= self.Cancel
            self.Form = None