#include "itkTimeProbe.h"
#include "itkBinaryThresholdImageFunction.h"
#include "itkImageLinearIteratorWithIndex.h"
#include "itkFloodFilledImageFunctionConditionalConstIterator.h"

#include "itkShapedFloodFilledImageFunctionConditionalConstIterator.h"

#include <fstream>
#include <sstream>

int main(int argc, char **argv)
{  
  const unsigned int ImageDimension = 2;
  typedef unsigned int PixelType;
  
  typedef itk::Image<PixelType, ImageDimension> ImageType;
  typedef ImageType::RegionType RegionType;
  typedef ImageType::IndexType IndexType;

  typedef itk::BinaryThresholdImageFunction<ImageType> FunctionType;
  typedef itk::FloodFilledImageFunctionConditionalConstIterator<ImageType, FunctionType> FloodFilledIteratorType;
  typedef itk::ShapedFloodFilledImageFunctionConditionalConstIterator<ImageType, FunctionType> ShapedFloodFilledIteratorType;  
  
  FunctionType::Pointer function = FunctionType::New();    
  
  static const unsigned int MAX_TIME_TESTS = 100;
  
  std::vector<IndexType> seedList;
  
  RegionType region;
  region.SetSize(0, 1280);
  region.SetSize(1, 1280);
  region.SetIndex(0, 0);
  region.SetIndex(1, 0);
  
  ImageType::Pointer inputImage = ImageType::New();
  inputImage->SetRegions(region);
  inputImage->Allocate();
  inputImage->FillBuffer(0);
  
  function->SetInputImage ( inputImage );  
  function->ThresholdAbove ( 1 ); // >= 1      
  
  itk::ImageLinearIteratorWithIndex<ImageType> it( inputImage, region );
  
  // make sure that we create a 4-connected image!
  for (int dir = 0; dir <= 1; ++dir)
  {
    it.SetDirection(dir);
    it.GoToBegin();
    while (!it.IsAtEnd())
    {
      while (!it.IsAtEndOfLine())
      {
        // add a seed
        if (seedList.empty()) seedList.push_back(it.GetIndex());
        
        it.Set(1);
        ++it;
      }
      
      // and jump over every 
      it.NextLine();
      if (!it.IsAtEnd()) it.NextLine();
      if (!it.IsAtEnd()) it.NextLine();
    }
  }  

  FloodFilledIteratorType floodIt(inputImage, function, seedList);
  ShapedFloodFilledIteratorType shapedFloodIt(inputImage, function, seedList);  
  shapedFloodIt.SetFullyConnected(false); // 4-connected, default

  std::vector<double> oldSamples;
  std::vector<double> newSamples;

  for (unsigned int j=0; j<MAX_TIME_TESTS; ++j)
  {
    itk::TimeProbe oldProbe, newProbe;
    
    int m,n;
    
    newProbe.Start();
    shapedFloodIt.GoToBegin();
    for (m = 0; !shapedFloodIt.IsAtEnd() ; ++shapedFloodIt, ++m)
    {
      const IndexType& index = shapedFloodIt.GetIndex();
    }
    newProbe.Stop();
    
    oldProbe.Start();
    floodIt.GoToBegin();
    for (n = 0; !floodIt.IsAtEnd() ; ++floodIt, ++n)
    {
      const IndexType& index = floodIt.GetIndex();
    }
    oldProbe.Stop();
    
    oldSamples.push_back(oldProbe.GetMeanTime());
    newSamples.push_back(newProbe.GetMeanTime());
    
    std::stringstream width, height;
    width << region.GetSize(0);
    height << region.GetSize(1);
    
    std::string filename = "timing" + width.str() + "x" + height.str() + ".txt";
    std::ofstream file(filename.c_str());
    
    std::vector<double>::const_iterator nIt = newSamples.begin();
    std::vector<double>::const_iterator oIt = oldSamples.begin();
    for (; nIt != newSamples.end(); ++nIt, ++oIt)
    {
      file << *oIt << "\t" << *nIt << std::endl;
    }
    
    file.close();
  }

  return 0;
}