#!/usr/bin/env python
# -*- coding: utf-8 -*-

#    Copyright 2010 C Sommer, C Straehle, U Koethe, FA Hamprecht. All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without modification, are
#    permitted provided that the following conditions are met:
#
#       1. Redistributions of source code must retain the above copyright notice, this list of
#          conditions and the following disclaimer.
#
#       2. Redistributions in binary form must reproduce the above copyright notice, this list
#          of conditions and the following disclaimer in the documentation and/or other materials
#          provided with the distribution.
#
#    THIS SOFTWARE IS PROVIDED BY THE ABOVE COPYRIGHT HOLDERS ``AS IS'' AND ANY EXPRESS OR IMPLIED
#    WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
#    FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE ABOVE COPYRIGHT HOLDERS OR
#    CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
#    CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
#    SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
#    ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#    NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
#    ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#    The views and conclusions contained in the software and documentation are those of the
#    authors and should not be interpreted as representing official policies, either expressed
#    or implied, of their employers.

"""
Watershed iterative segmentation plugin
"""


import vigra, numpy

from segmentorBase import *
import traceback
from enthought.traits.api import *
from enthought.traits.ui.api import *

ok = False

try:
    import vigra.svs
    ok = True
except Exception, e:
    print "####################################"
    print
    print "    please update to the latest"
    print "    private viga repo !!!"
    print
    print "####################################"

if ok:
    class SegmentorWSiter(SegmentorBase):
        name = "Supervoxel Segmentation"
        description = "Segmentation plug-in using sparse basin graph"
        author = "HCI, University of Heidelberg"
        homepage = "http://hci.iwr.uni-heidelberg.de"

        showBorders = Bool(False)
        edgeWeights = Enum("Average", "Difference")
        algorithm = Enum("Graphcut", "Watershed")        
        bias = Float(0.95)
        biasedLabel = Int(1)
        sigma = Float(0.2)
        
        view = View( Item('showBorders'), Item('edgeWeights'), Item('algorithm'), Item('bias'),  Item('biasedLabel'), Item('sigma'), buttons = ['OK', 'Cancel'],  )        
        
        lastBorderState = False        
        
        class IndexedAccessor:
            """
            Helper class that behaves like an ndarray, but does a Lookuptable access
            """

            def __init__(self, volumeBasins, basinLabels):
                self.volumeBasins = volumeBasins
                self.basinLabels = basinLabels
                self.dtype = basinLabels.dtype
                self.shape = volumeBasins.shape

            def __getitem__(self, key):
                return self.basinLabels[self.volumeBasins[tuple(key)]]

            def __setitem__(self, key, data):
                #self.data[tuple(key)] = data
                print "##########ERROR ######### : SegmentationDataAccessor setitem should not be called"

        def segment3D(self, labelVolume, labelValues, labelIndices):
            print "setting seeds"
            self.segmentor.setSeeds(labelValues, labelIndices)
            if self.algorithm == "Graphcut":
                print "Executing Graphcut with sigma = %d"  % (self.sigma,)
                self.basinLabels = self.segmentor.doGC(self.sigma)
            elif self.algorithm == "Watershed":
                print "Executing Watershed with bias %d and biasedLabel %d" % (self.bias,  self.biasedLabel,)
                self.basinLabels = self.segmentor.doWS(self.bias,  self.biasedLabel)

            if self.lastBorderState != self.showBorders:
                self.getBasins()
            
            self.acc = SegmentorWSiter.IndexedAccessor(self.volumeBasins, self.basinLabels)
            return self.acc

        def segment2D(self, labels):
            pass


        def setupWeights(self, weights):
            self.weights = weights
            #self.weights = numpy.average(weights, axis = 3).astype(numpy.uint8)#.swapaxes(0,2).view(vigra.ScalarVolume)#
            if weights.dtype != numpy.uint8:
                print "converting weights to uint8"
                self.weights = weights.astype(numpy.uint8)
                
#            self.weights = numpy.zeros(weights.shape[0:-1], 'uint8')
#            self.weights[:] = 3
#            self.weights[:,:,0::4] = 10
#            self.weights[:,0::4,:] = 10
#            self.weights[0::4,:,:] = 10
#            self.weights = self.weights

            #self.ws = vigra.tws.IncrementalWS(self.weights, 0)
            print "Incoming weights :", self.weights.dtype, self.weights.shape

            if hasattr(self, "segmentor"):
                del self.segmentor
                del self.volumeBasins
            

            if self.edgeWeights == "Difference":
                useDifference = True
            else:                
                useDifference = False
            
            self.segmentor = vigra.svs.segmentor(self.weights, useDifference, 0, 255, 2048)
            
            self.getBasins()
            
            print "Outgoing weights :", self.volumeBasins.shape

            self.volumeBasins.shape = self.volumeBasins.shape + (1,)

        def getBasins(self):
            self.lastBorderState = self.showBorders
            if self.showBorders:
                self.volumeBasins = self.segmentor.getVolumeBasinsWithBorders() #WithBorders()
            else:
                self.volumeBasins = self.segmentor.getVolumeBasins()            
@@ -2536,16 +2536,29 @@
  = Enum(
+%22Watershed%22, 
 %22Graphcu
@@ -2562,25 +2562,27 @@
 hcut%22, %22
-Watershed
+Random Walk
 %22)      
@@ -2667,16 +2667,72 @@
 at(0.2)%0A
+        lis_options = String(%22-i bicgstab -tol 1.0e-9%22)%0A
         
@@ -2865,16 +2865,36 @@
 igma'), 
+Item('lis_options'),
 buttons 
@@ -4289,16 +4289,284 @@
 edLabel)
+%0A            elif self.algorithm == %22Random Walk%22:%0A                print %22Executing Random Walk with sigma %25f, and lis options %25s%22 %25 (self.sigma,  self.lis_options,)%0A                self.basinLabels = self.segmentor.doRW(self.sigma,  self.lis_options)%0A                
 %0A%0A      

add supervoxel random walker support