//
// This file is part of MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2015 Richard S. Smith and collaborators.
//
// If you use MorphoGraphX in your work, please cite:
//   http://dx.doi.org/10.7554/eLife.05864
//
// MorphoGraphX is free software, and is licensed under under the terms of the 
// GNU General (GPL) Public License version 2.0, http://www.gnu.org/licenses.
// 
#ifndef ITKSEGMENTATION_HPP
#define ITKSEGMENTATION_HPP

#include <ITKConfig.hpp>
#include <ITKProcess.hpp>

namespace mgx 
{
  class mgxITK_EXPORT ITKWatershed : public Process {
  public:
    ITKWatershed(const Process& process) : Process(process) 
    {
      setName("Stack/ITK/Segmentation/ITK Watershed");
      setDesc("Seeded watershed segmentation, image in main store, labels in work");
      setIcon(QIcon(":/images/SegmentMesh.png"));

      addParm("Fully Connected", "Fully connected", "No", booleanChoice());
      addParm("Mark Watershed Line", "Mark Watershed Line", "No", booleanChoice());	
	  }
  
    bool run()
    {
      if(!checkState().store(STORE_MAIN | STORE_NON_LABEL)
                            .store(STORE_WORK | STORE_LABEL))
        return false;
      Stack* stack = currentStack();
      Store* main = stack->main();
      Store* labels = stack->work();
      bool connect8 = stringToBool(parm("Fully connected"));
      bool markLine = stringToBool(parm("Mark Watershed Line"));
      return run(main, labels, connect8, markLine);
    }
  
    bool run(const Store* image, Store* labels, bool connect8, bool markLine);
  };
  
  class mgxITK_EXPORT ITKWatershedAutoSeeded : public Process {
  public:
    ITKWatershedAutoSeeded(const Process& process) : Process(process) 
    {
	  setName("Stack/ITK/Segmentation/ITK Watershed Auto Seeded");
	  setDesc("Auto Seeded Watershed segmentation");
	  setIcon(QIcon(":/images/SegmentMesh.png"));

	  addParm("Fully connected","Fully connected","No",booleanChoice());
	  addParm("Mark Watershed Line","Mark Watershed Line","No",booleanChoice());
	  addParm("Level","Level","1500");						
	}
  
    bool run()
    {
      if(!checkState().store(STORE_VISIBLE | STORE_NON_LABEL))
        return false;
      Stack* stack = currentStack();
      Store* main = stack->currentStore();
      Store* labels = stack->work();
      bool connect8 = stringToBool(parm("Fully connected"));
      bool markLine = stringToBool(parm("Mark Watershed Line"));
      if(run(stack, main, labels, parm("Level").toUInt(), connect8, markLine)) {
        main->hide();
        labels->show();
        return true;
      }
      return false;
    }
  
    bool run(Stack* stack, const Store* image, Store* labels, ushort level, 
                                                               bool connect8, bool markLine);
  
  };
}

#endif
