//
// 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.
// 
#include "ITKProgress.hpp"

#include "ITKSegmentation.hpp"

#include <itkMorphologicalWatershedFromMarkersImageFilter.h>
#include <itkMorphologicalWatershedImageFilter.h>

namespace mgx 
{
  bool ITKWatershed::run(const Store* image, Store* labels, bool connect8, bool markLine)
  {
    UImageConverter::Pointer converterImage = UImageConverter::New();
    UImageConverter::Pointer converterLabel = UImageConverter::New();
  
    converterImage->SetStore(image);
    converterLabel->SetStore(labels);
  
    typedef itk::MorphologicalWatershedFromMarkersImageFilter<UImageType, UImageType> FilterType;
  
    FilterType::Pointer filter = FilterType::New();
  
    filter->SetFullyConnected(connect8);
    filter->SetMarkWatershedLine(markLine);
    filter->SetInput1(converterImage->GetOutput());
    filter->SetInput2(converterLabel->GetOutput());
  
    ITKProgress progress("Seeded Watershed");
    progress.setFilter(filter);
  
    filter->Update();
  
    UImageType::ConstPointer result = filter->GetOutput();
  
    if(!UImageConverter::TransferImage(labels->stack(), labels, result)) {
      setErrorMessage("ITK Process tried to change the size of the stack");
      return false;
    }
  
    labels->changed();
  
    return true;
  }
  
  REGISTER_PROCESS(ITKWatershed);
  
  bool ITKWatershedAutoSeeded::run(Stack* stack, const Store* image, Store* labels, ushort level, bool connect8, bool markLine)
  {
    UImageConverter::Pointer converterImage = UImageConverter::New();
    UImageConverter::Pointer converterLabel = UImageConverter::New();
  
    converterImage->SetStore(image);
    converterLabel->SetStore(labels);
  
    typedef itk::MorphologicalWatershedImageFilter<UImageType, UImageType> FilterType;
  
    FilterType::Pointer filter = FilterType::New();
  
    filter->SetLevel(level);
    filter->SetFullyConnected(connect8);
    filter->SetMarkWatershedLine(markLine);
    filter->SetInput(converterImage->GetOutput());
  
    ITKProgress progress("Auto-Seeded Watershed");
    progress.setFilter(filter);
  
    filter->Update();
    if(filter->GetAbortGenerateData())
      return false;
  
    UImageType::ConstPointer result = filter->GetOutput();
  
    if(!UImageConverter::TransferImage(labels->stack(), labels, result)) {
      setErrorMessage("ITK Process tried to change the size of the stack");
      return false;
    }
  
    labels->setLabels(true);
  
    setCurrentStackId(stack->id());
    stack->setCurrentStore(labels);
    if(!runProcess("Stack/Segmentation/Relabel", QStringList() << "1" << "1"))
      return false;
  
    labels->changed();
  
    return true;
  }
  REGISTER_PROCESS(ITKWatershedAutoSeeded);
}
