# Avizo-Script-Object V3.0
# NO DEPRECATION WARNING

##############################################################
#  Script for FIB processing
#
#    Performs operations for FIB processing (shearing, cropping, alignment, shadow correction)
#    Can be used as a basis for advanced processing of FIB images.
#    As it is based on ActionStep2.scro in order to benefit from action based workflow,
#    user interaction is reduced
#
#
#    Workflow:
#    - optional croping before alignment
#        Is often used to restrict the volume data to the area of interest
#    - alignment
#        Internally uses Avizo Align module. Rotation can be disabled. A label volume data defining
#        the alignment area can be connected to the label port of the script object.
#    - optional croping after alignment
#    - optional shearing correction
#    - shadow correction
#
# author: Gwenole Tallec (gtallec@fei.com)
# (C) copyright 2011, FEI SAS
#
##############################################################

# #########################################################################
# Load ActionStep2.scro file which contains some definition of procedures
# #########################################################################
set file share/script-objects/ActionStep2.scro
if [file readable $AMIRA_LOCAL/$file] {
    source $AMIRA_LOCAL/$file
} else {
    source $AMIRA_ROOT/$file
}

# #########################################################################
# Destructor
# #########################################################################
"$this" proc destructor {} {

    # Close crop editor if exists
    if {[exists cropEditor]} {
        cropEditor detach
    }

    # Remove modules
    set mainOrthoSlice       ["$this" getVar mainOrthoSlice]
    set alignedSliceModule   ["$this" getVar alignedSliceModule]
    set maskThresholdModule  ["$this" getVar maskThresholdModule]
    set shearModule          ["$this" getVar shearModule]

    remove "$alignedSliceModule"
    remove "$maskThresholdModule"
    remove "$shearModule"
    remove "$mainOrthoSlice"

    # Retrieve the initial state of the camera
    viewer setCameraType ["$this" getVar initialCameraType]

    # Remove all intermediate data but result
    set data ["$this" getVar currentData]
    set resultData ["$this" getVar resultData]

    # We disconnect all the data that are inputs for the script.
    # This aims at protecting the script against unwanted deletion.
    # Indeed, if we delete an input of the module the module is deleted.
    "$this" labelForAlignment disconnect
    "$this" mask disconnect
    "$this" croppedData disconnect

    if { ["$this" getVar dupDataLabel] != "" } {
        set dataToDelete ["$this" getVar dupDataLabel]
        if { ("$dataToDelete" != "$data") && ("$dataToDelete" != "$resultData") } {
            remove "$dataToDelete"
        }
    }

    if { ["$this" getVar alignedData] != "" } {
        set dataToDelete ["$this" getVar alignedData]
        if { ("$dataToDelete" != "$data") && ("$dataToDelete" != "$resultData") } {
            remove "$dataToDelete"
        }
    }

    if { ["$this" getVar shearedData] != "" } {
        set dataToDelete ["$this" getVar shearedData]
        if { ("$dataToDelete" != "$data") && ("$dataToDelete" != "$resultData") } {
            remove "$dataToDelete"
        }
    }

    if { ["$this" getVar shearedCorrectedData] != "" } {
        set dataToDelete ["$this" getVar shearedCorrectedData]
        if { ("$dataToDelete" != "$data") && ("$dataToDelete" != "$resultData") } {
            remove "$dataToDelete"
        }
    }

    if { ["$this" getVar alignedCroppedData] != "" } {
        set dataToDelete ["$this" getVar alignedCroppedData]
        if { ("$dataToDelete" != "$data") && ("$dataToDelete" != "$resultData") } {
            remove "$dataToDelete"
        }
    }
}


# #########################################################################
# Define action if the OK button of Crop Editor is hit
# #########################################################################
"$this" proc compute {} {
    if {!["$this" getVar cropApply]} {
        if {[exists cropEditor]} {
            set initialDims ["$this" getVar initialDims]
            set croppedData ["$this" croppedData source]
            set actualDims [$croppedData getDims]

            if {[lindex $initialDims 0] != [lindex $actualDims 0]
            || [lindex $initialDims 1] != [lindex $actualDims 1]
            || [lindex $initialDims 2] != [lindex $actualDims 2]} {
                "$this" croppedData disconnect
                "$this" setVar cropCompute 1
                eval ["$this" action getCmd 2]
            }
        }
    }
}

# #########################################################################
# Define all action of this wizard
# #########################################################################
"$this" proc defineActions {} {

    # resultData will be the data supposed to be the result, thus modified at each step
    "$this" setVar resultData ""

    # Port used to give a label for alignment
    "$this" newPortConnection labelForAlignment HxUniformLabelField3
    "$this" labelForAlignment setLabel "Label For Alignment"
    "$this" labelForAlignment hide

    # Port used to give a label for alignment
    "$this" newPortConnection mask HxUniformLabelField3
    "$this" mask setLabel "Mask"

    # Save the initial state of the camera
    "$this" setVar initialCameraType [viewer getCameraType]

    # Switch camera to Orthographic
    viewer setCameraType Orthographic

    "$this" newPortConnection croppedData HxUniformScalarField3
    "$this" croppedData hide

    # Now declare the actions

    # ######################### #
    # CROPPING BEFORE ALIGNMENT #
    # ######################### #
    "$this" addAction "Crop Before Alignment" "Crop before alignment" initAction {

        "$this" setVar cropApply 0
        "$this" setVar cropCompute 0

        # Duplicate input because we will need cropping
        # If network is reloading, duplicated data has already been created and saved
        if { !["$this" getVar isNetworkLoading] } {
            set currentData ["$this" getVar currentData]
            # Disconnect MainOrthoSlice
            ["$this" getVar mainOrthoSlice] data disconnect

            # Cleanup
            if { ["$this" getVar dupDataLabel] != "" } {
                remove ["$this" getVar dupDataLabel]
            }

            # Duplicate data
            set dupData [$currentData duplicate]
            # If dupDataLabel has not been initialize, compose it (based on data's label)
            if { ["$this" getVar dupDataLabel] == "" } {
                "$this" setVar dupDataLabel [$currentData getLabel]-dup
            }

            # Data can be automatically renamed if dupDataLabel is already used
            "$this" setVar dupDataLabel [$dupData setLabel ["$this" getVar dupDataLabel]]
        }

        set dupDataLabel ["$this" getVar dupDataLabel]

        # Connect MainOrthoSlice to duplicated data
        "$this" attachMainOrthoSlice $dupDataLabel

        # Attach a crop editor to duplicated data
        create HxImageCrop
        cropEditor attach $dupDataLabel
        "$this" croppedData connect $dupDataLabel
        "$this" setVar initialDims [$dupDataLabel getDims]

    } onSkipAction {

        "$this" croppedData disconnect
        "$this" setVar cropApply 1
        # Detach if user didn't press OK or Cancel
        if { [exists cropEditor] } {
            cropEditor detach
        }

    } onApplyAction {

        if {!["$this" getVar cropCompute]} {
            "$this" croppedData disconnect
            "$this" setVar cropApply 1

            # Crop if user didn't press OK or Cancel
            if { [exists cropEditor] } {
                cropEditor doIt
            }
        }
    }

    # ######### #
    # ALIGNMENT #
    # ######### #
    "$this" addAction "Align Slices" "Alignment" initPorts {

        # Create alignment options
        "$this" newPortRadioBox alignMethod 2
        "$this" alignMethod setLabel "Align Method"
        "$this" alignMethod setTooltip "Defines the method used for the alignment ( gravity centers or least-squares)"
        "$this" alignMethod hide
        "$this" alignMethod setLabel 0 "Gravity centers"
        "$this" alignMethod setLabel 1 "Least-squares"
        "$this" alignMethod setValue 1
        "$this" alignMethod untouch

        # Create alignment options
        "$this" newPortToggleList alignOptions 1
        "$this" alignOptions setLabel "Align Options"
        "$this" alignOptions setTooltip "Enable/Disable rotations between 2 consecutive slices"
        "$this" alignOptions hide
        "$this" alignOptions setLabel "Enable rotations"
        "$this" alignOptions setValue 0 0
        "$this" alignOptions untouch

    } initAction {

        set dupDataLabel ["$this" getVar dupDataLabel]
        "$this" setVar resultData $dupDataLabel

        # Connect MainOrthoSlice to duplicated data
        "$this" attachMainOrthoSlice $dupDataLabel

        # Cleanup
        if { ["$this" getVar alignedData] != $dupDataLabel} {
            remove ["$this" getVar alignedData]
        }
        remove ["$this" getVar alignedCroppedData]

        # Create the AlignSlices module if necessary
        if { ![exists ["$this" getVar alignedSliceModule]] } {
            set alignedSliceModule [create HxAlignSlices]
            "$this" setVar alignedSliceModule $alignedSliceModule
        }

        set alignedSliceModule ["$this" getVar alignedSliceModule]

        $alignedSliceModule data connect ["$this" getVar dupDataLabel]

        if { [exists ["$this" mask source]] } {
            $alignedSliceModule mask connect ["$this" mask source]
        }
        $alignedSliceModule fire

    } onSkipAction {

        "$this" setVar alignedData ["$this" getVar dupDataLabel]

    } onApplyAction {

        set alignMode ["$this" alignMethod getValue]
        set alignedSliceModule ["$this" getVar alignedSliceModule]

        # Check rotations enabled
        $alignedSliceModule setRotationAllowed ["$this" alignOptions getValue 0]

        # Needed before computation : open editor
        $alignedSliceModule action hit 0
        $alignedSliceModule fire

        $alignedSliceModule setReferenceSlice 0
        $alignedSliceModule setEditableSlice upper

        # Check method
        $alignedSliceModule setAlignMode ["$this" alignMethod getValue]
        $alignedSliceModule fire

        # Perform alignment
        $alignedSliceModule alignAll

        # Resample result
        $alignedSliceModule action hit 1
        $alignedSliceModule fire

        # Close editor
        $alignedSliceModule save
        $alignedSliceModule action hit 0
        $alignedSliceModule fire

        "$this" setVar alignedData [$alignedSliceModule getResult]

        # Disconnect portMaster of alignedData in order to correctly reload saved project
        ["$this" getVar alignedData] master disconnect

    }

    # #################### #
    # CROP AFTER ALIGNMENT #
    # #################### #
    "$this" addAction "Crop After Alignment" "Crop after alignment" initAction {

        "$this" setVar cropApply 0
        "$this" setVar cropCompute 0

        # connect MainOrthoSlice
        "$this" attachMainOrthoSlice ["$this" getVar alignedData]

        # Duplicate input because we will need cropping
        if { !["$this" getVar isNetworkLoading] } {
            "$this" setVar alignedCroppedData ["$this" getVar alignedData]-cropped
        }
        set alignedCropedData ["$this" getVar alignedCroppedData]

        remove $alignedCropedData

        set dupData [["$this" getVar alignedData] duplicate]
        $dupData setLabel $alignedCropedData

        # Create crop
        create HxImageCrop
        cropEditor attach $alignedCropedData

        "$this" setVar resultData $alignedCropedData

        "$this" croppedData connect $alignedCropedData
        "$this" setVar initialDims [$alignedCropedData getDims]

    } onSkipAction {

        "$this" croppedData disconnect
        "$this" setVar cropApply 1
        # Detach if user didn't press OK or Cancel
        if {[exists cropEditor]} {
            cropEditor detach
        }

    } onBackAction {

        "$this" croppedData disconnect
        "$this" setVar cropApply 1
        # Detach if user didn't press OK or Cancel
        if {[exists cropEditor]} {
            cropEditor detach
        }

    } onApplyAction {

        if {!["$this" getVar cropCompute]} {
            "$this" croppedData disconnect
            "$this" setVar cropApply 1

            # Crop if user didn't press OK or Cancel
            if { [exists cropEditor] } {
                cropEditor doIt
            }
        }
    }

    # ################### #
    # SHEARING CORRECTION #
    # ################### #
    "$this" addAction Shearing "Shearing correction" initPorts {

        "$this" newPortFloatTextN shearAngle 1
        "$this" shearAngle setLabel "Shear Angle"
        "$this" shearAngle setTooltip "Angle used for shearing correction."
        "$this" shearAngle hide
        "$this" shearAngle setValue 56
        "$this" shearAngle untouch

    } initAction {

        set alignedCroppedData ["$this" getVar alignedCroppedData]

        # Remove previously sheared data
        if { (["$this" getVar shearedData] != $alignedCroppedData) && !["$this" getVar isNetworkLoading]} {
            ["$this" getVar mainOrthoSlice] data disconnect
            remove ["$this" getVar shearedData]
        }

        # Connect MainOrthoSlice
        "$this" attachMainOrthoSlice $alignedCroppedData

        "$this" setVar resultData $alignedCroppedData

    } onSkipAction {

        set alignedCroppedData ["$this" getVar alignedCroppedData]
        "$this" setVar shearedData $alignedCroppedData

    } onApplyAction {

        # Create the Shearing correction module
        if {![exists ["$this" getVar shearModule]]} {
            set shearModule [create HxShear {shear}]
            "$this" setVar shearModule "$shearModule"
        }

        set alignedCroppedData ["$this" getVar alignedCroppedData]
        set shearModule        ["$this" getVar shearModule]

        "$shearModule" angle setValue ["$this" shearAngle getValue 0]
        "$shearModule" data connect $alignedCroppedData
        "$shearModule" action hit
        "$shearModule" fire

        "$this" setVar shearedData ["$shearModule" getResult]
        ["$this" getVar shearedData] master disconnect
    }

    # ################## #
    # SHADING CORRECTION #
    # ################## #
    "$this" addAction "Shading Correction" "Shading correction" initPorts {

        "$this" newPortRangeSlider maskThreshold
        "$this" maskThreshold setLabel "Mask Threshold"
        "$this" maskThreshold setTooltip "Threshold defining the background used for shading correction"

        "$this" newPortFloatTextN shadingCorrection 1
        "$this" shadingCorrection setLabel "Shading Correction"
        "$this" shadingCorrection setValue 200

    } initAction {

        set shearedData ["$this" getVar shearedData]

        # Connect MainOrthoSlice
        "$this" attachMainOrthoSlice "$shearedData"

        if { ![exists ["$this" getVar maskThresholdModule]] } {
            set maskThresholdModule [create HxInteractiveThreshold]

            "$this" setVar maskThresholdModule "$maskThresholdModule"
            "$maskThresholdModule" hideIcon

            # Interconnect threshold and orientation
            "$maskThresholdModule" threshold connect "$this" maskThreshold
            "$maskThresholdModule" orientation  connect "$this" sliceOrientation
        }

        set maskThresholdModule ["$this" getVar maskThresholdModule]

        "$maskThresholdModule" data connect "$shearedData"
        "$maskThresholdModule" sliceNumber connect "$this" sliceNumber
        "$maskThresholdModule" rendering setValue 0
        "$maskThresholdModule" fire
        "$maskThresholdModule" setViewerMask 1
        "$maskThresholdModule" fire

        "$this" setVar resultData "$shearedData"

        # Close connection editor
        "$this" closeConnectionEditor

    } onSkipAction {

        # Hide threshold stuffs
        set maskThresholdModule ["$this" getVar maskThresholdModule]
        if { [exists "$maskThresholdModule"] } {
            "$maskThresholdModule" setViewerMask 0
        }

        "$this" setVar shearedCorrectedData ["$this" getVar shearedData]

    } onBackAction {

        set maskThresholdModule ["$this" getVar maskThresholdModule]

        # Hide threshold stuffs
        if {[exists "$maskThresholdModule"]} {
            "$maskThresholdModule" setViewerMask 0
        }

    } onApplyAction {

        set maskThresholdModule ["$this" getVar maskThresholdModule]

        "$maskThresholdModule" doIt hit
        "$maskThresholdModule" fire
        "$maskThresholdModule" setViewerMask 0

        set shearedData ["$this" getVar shearedData]
        set shearedThresholdedData ["$maskThresholdModule" getResult]

        if { ["$this" getVar shearedCorrectedData] == "" } {
            "$this" setVar shearedCorrectedData "$shearedData-corrected"
        }

        set shearedCorrectedData ["$this" getVar shearedCorrectedData]

        set bkgimgModule [moduleApply bkgimg inputImage "$shearedData" inputImageMask "$shearedThresholdedData"]
        set bkg [$bkgimgModule getResult]
        remove $bkgimgModule
        remove "$shearedThresholdedData"

        set shadingCorrectionParam ["$this" shadingCorrection getValue]
        set shading_corimageModule [moduleApply shading_corimage inputImage "$shearedData" whiteReference "$bkg" normalizationFactor $shadingCorrectionParam resultSlot "$shearedCorrectedData"]
        set shearedCorrectedData [$shading_corimageModule getResult]
        remove $shading_corimageModule
        remove $bkg

        "$this" setVar shearedCorrectedData "$shearedCorrectedData"

        # Connect MainOrthoSlice
        "$this" attachMainOrthoSlice "$shearedCorrectedData"

        "$this" setVar resultData "$shearedCorrectedData"
    }
}

# #########################################################################
# Initialize variables
# #########################################################################
"$this" proc initVar {} {

    # Initialize module variables
    if { [lsearch ["$this" getVarNames] alignedSliceModule]     == -1} {"$this" setVar alignedSliceModule     {}}
    if { [lsearch ["$this" getVarNames] maskThresholdModule]    == -1} {"$this" setVar maskThresholdModule    {}}
    if { [lsearch ["$this" getVarNames] shearModule]            == -1} {"$this" setVar shearModule            {}}

    # Initialize data variables
    if { [lsearch ["$this" getVarNames] resultData]             == -1} {"$this" setVar resultData             {}}
    if { [lsearch ["$this" getVarNames] dupDataLabel]           == -1} {"$this" setVar dupDataLabel           {}}
    if { [lsearch ["$this" getVarNames] alignedData]            == -1} {"$this" setVar alignedData            {}}
    if { [lsearch ["$this" getVarNames] shearedData]            == -1} {"$this" setVar shearedData            {}}
    if { [lsearch ["$this" getVarNames] shearedCorrectedData]   == -1} {"$this" setVar shearedCorrectedData   {}}
    if { [lsearch ["$this" getVarNames] alignedCroppedData]     == -1} {"$this" setVar alignedCroppedData     {}}
}