/*=========================================================================

  Program:   Contour Interpolation with Geodesic Snakes
  Module:    $RCSfile: GeoInterp.cxx,v $
  Language:  C++
  Date:      $Date: 2006/06/30 10:00:00 $
  Version:   $Revision: 1.0 $

  Copyright (c) 2003 Insight Software Consortium
  All rights reserved.

  Redistribution and use in source and binary forms, with or without
  modification, are permitted provided that the following conditions are met:

   * Redistributions of source code must retain the above copyright notice,
     this list of conditions and the following disclaimer.
   * 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.
   * The name of the Insight Consortium, nor the names of any consortium members,
     nor of any contributors, may be used to endorse or promote products derived
     from this software without specific prior written permission.
   * Modified source versions must be plainly marked as such, and must not be
     misrepresented as being the original software.

 THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER AND CONTRIBUTORS ``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 AUTHORS 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.

=========================================================================*/


#include "CommandLineArgumentParser.h"
#include "ImageCoordinateGeometry.h"
#include "ImageIORoutines.h"
#include "IRISApplication.h"
#include "IRISException.h"
#include "IRISImageData.h"
#include "SNAPRegistryIO.h"
#include "SystemInterface.h"

#include "itkImage.h"
#include "itkNumericTraits.h"

#include <itksys/SystemTools.hxx>

#include "Registry.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkRawImageIO.h"
#include "SNAPImageData.h"
#include "IRISMeshPipeline.h"
#include "vtkPolyData.h"
#include "vtkBYUWriter.h"
#include "vtkTriangleFilter.h"
#include "SNAPLevelSetDriver.h"

using namespace std;


/** \class Automate<TPixel>
 * \brief Automates Geodesic Snakes driver
 */
template < class TPixel > class Automate
{
  // Image type definition
  typedef itk::Image < TPixel, 3 > ImageType;
  typedef itk::Image < float, 3 > FloatImageType;
  typedef itk::Image < LabelType, 3 > LabelImageType;
  typedef typename itk::SmartPointer < ImageType > ImagePointer;

  // Image IO type definition
  typedef itk::ImageIOBase ImageIOType;
  typedef typename itk::SmartPointer < ImageIOType > ImageIOPointer;

public:
  /**
   * Saves the image passed in saveImg to a file named fname.
   */
  template < class SaveImageType >
  void saveImage (SaveImageType * saveImg, const char *fname)
  {
    typedef typename itk::ImageFileWriter < SaveImageType > Writer;
    typedef typename Writer::Pointer WriterPointer;

    WriterPointer writer = Writer::New ();
    ImageIOPointer writerIO =
      CreateRawImageIO (typename SaveImageType::PixelType (), 0);

    typename SaveImageType::SizeType imgSize =
      saveImg->GetLargestPossibleRegion ().GetSize ();
    writerIO->SetDimensions (0, imgSize[0]);
    writerIO->SetDimensions (1, imgSize[1]);
    writerIO->SetDimensions (2, imgSize[2]);
    cout << "Saving image <" << fname << "> [" << imgSize[0] << "," <<
      imgSize[1] << "," << imgSize[2] << "] ...";
    writer->SetImageIO (writerIO);
    writer->SetFileName (fname);

    writer->SetInput (saveImg);
    writer->Update ();

    cout << " done" << endl;
  }

  /**
   * Loads the image: the filename is obtained from the registry.
   */
  bool LoadImage ()
  {
    bool rc = true;
		const unsigned int header = (unsigned int) params["HeaderLength"][0];

		switch ((int) params["VoxelType"][0])
			{
			case 0:
				m_ImageIO = CreateRawImageIO ((unsigned char) 0, header);
				break;
			case 1:
				m_ImageIO = CreateRawImageIO ((char) 0, header);
				break;
			case 2:
				m_ImageIO = CreateRawImageIO ((unsigned short) 0, header);
				break;
			case 3:
				m_ImageIO = CreateRawImageIO ((short) 0, header);
				break;
			case 4:
				m_ImageIO = CreateRawImageIO ((unsigned int) 0, header);
				break;
			case 5:
				m_ImageIO = CreateRawImageIO ((int) 0, header);
				break;
			case 6:
				m_ImageIO = CreateRawImageIO (0.0f, header);
				break;
			default:
				assert (0);
			}


    // Try to load a file
    try
    {
      typedef itk::ImageFileReader < ImageType > ReaderType;
      typedef typename ReaderType::Pointer ReaderPointer;

      // Set the dimensions of the image
      m_ImageIO->SetDimensions (0, (unsigned int) params["DimX"][0]);
      m_ImageIO->SetDimensions (1, (unsigned int) params["DimY"][0]);
      m_ImageIO->SetDimensions (2, (unsigned int) params["DimZ"][0]);

      // Set the spacing of the image
      m_ImageIO->SetSpacing (0, params["SpacingX"][0.0]);
      m_ImageIO->SetSpacing (1, params["SpacingY"][0.0]);
      m_ImageIO->SetSpacing (2, params["SpacingZ"][0.0]);

      // Set the endianness
      if (params["ByteAlign"][false] == true)
        m_ImageIO->SetByteOrderToLittleEndian ();
      else
        m_ImageIO->SetByteOrderToBigEndian ();

      // Create a reader
      ReaderPointer reader = ReaderType::New ();

      cout << "Loading ... ";

      // Set the file name
      reader->SetFileName (params["GreyImage"]["no such file"]);

      reader->SetImageIO (m_ImageIO);

      // Perform the read
      reader->Update ();

      // Get the output
      m_Image = reader->GetOutput ();

      // Disconnect the image from the reader
      m_Image->DisconnectPipeline ();

      // Store the ImageIO
      m_ImageIO = reader->GetImageIO ();

      cout << "done" << endl;
    }
    catch (itk::ExceptionObject & exc)
    {
      // Clear the image and image IO
      m_Image = NULL;
      m_ImageIO = NULL;

      // Show the error
      cerr << "Error reading image: " << exc.GetDescription () << endl;
      rc = false;
    }

    // Check if the image is valid (subclasses can perform extra tasks here)
    return rc;
  }

  /**
   * Creates an instance of RawImageIO with pixel type of the first argument
   * and header size as header.
   */
  template < class TRawPixel >
  itk::ImageIOBase::Pointer
  CreateRawImageIO (TRawPixel itkNotUsed (dummy), unsigned int header)
  {
    typedef itk::RawImageIO < TRawPixel, 3 > IOType;

    typename IOType::Pointer rawIO = IOType::New ();

    rawIO->SetHeaderSize (header);

    itk::ImageIOBase::Pointer baseIO = rawIO.GetPointer ();

    return baseIO;
  }

  /**
   * Loads all the parameters from the command file and runs the contour
   * interpolation algorithm. Takes an optional override for the image file.
   */
  void RunSegmentation (const char *filename, const char* greyImage = NULL)
  {
    //
    // Load the parameters file
    //
    cout << "Loading registry ... ";
    params = Registry (filename);
	if( greyImage != NULL )
	  {
	  cout << "Using image " << greyImage << " as input." << endl;
	  params["GreyImage"]	= string(greyImage);
	  }
    cout << "done" << endl;

    debuggingLevel = (unsigned int) params["DebuggingLevel"][0];

    //
    // Load the image
    //
    if (!LoadImage ())
      {
      cerr << "Error loading image" << endl;
      exit (0);
      }

    //
    // Preprocess the image
    //
    cout << "Preprocessing the image (filling with don't cares) ";
    typename ImageType::SizeType size =
      m_Image->GetLargestPossibleRegion ().GetSize ();
    const int xi   = 0;
    const int yi   = 1;
    const int zi   = 2;
    const int xdim = size[xi];
    const int ydim = size[yi];
    const int zdim = size[zi];

    FloatImageType::Pointer img = FloatImageType::New ();
    img->SetRegions (size);
    img->SetSpacing (m_Image->GetSpacing ());
    img->Allocate ();

    typename ImageType::IndexType idxLower, idxUpper;

    {
      const double IN_INTENSITY        = +1.0;
      const double OUT_INTENSITY       = -1.0;
      const double DONT_CARE_INTENSITY = 0.0;

      typename ImageType::IndexType i;

      idxLower[xi] = xdim;
      idxLower[yi] = ydim;
      idxLower[zi] = zdim;
      idxUpper[xi] = 0;
      idxUpper[yi] = 0;
      idxUpper[zi] = 0;
      int zmin = zdim;
      int zmax = 0;

      // first find start zslice and end zslice and identify the outer regions.
      for (i[zi] = 0; i[zi] != zdim && zmin == zdim; i[zi]++)
        {
        for (i[xi] = 0; i[xi] != xdim; i[xi]++)
          {
          for (i[yi] = 0; i[yi] != ydim; i[yi]++)
            {
            const TPixel p = m_Image->GetPixel (i);

            if (p != 0)
              {
              zmin = i[zi];
              idxLower[zi] = zmin;
              img->SetPixel (i, IN_INTENSITY);
              }
            else
              {
              img->SetPixel (i, OUT_INTENSITY);
              }
            }
          }
        }
      cout << '.';
      for (i[zi] = zdim - 1; i[zi] >= 0 && zmax == 0; i[zi]--)
        {
        for (i[xi] = 0; i[xi] != xdim; i[xi]++)
          {
          for (i[yi] = 0; i[yi] != ydim; i[yi]++)
            {
            const TPixel p = m_Image->GetPixel (i);

            if (p != 0)
              {
              zmax = i[zi];
              idxUpper[zi] = zmax;
              img->SetPixel (i, IN_INTENSITY);
              }
            else
              {
              img->SetPixel (i, OUT_INTENSITY);
              }
            }
          }
        }
      cout << '.';

      // Step from zmin+1 to zmax-1 filling all non-contoured slices with don't cares.
      for (i[zi] = zmin + 1; i[zi] <= zmax - 1; i[zi]++)
        {
        bool contoured = false;

        for (i[xi] = 0; i[xi] != xdim; i[xi]++)
          {
          for (i[yi] = 0; i[yi] != ydim; i[yi]++)
            {
            const TPixel p = m_Image->GetPixel (i);

            if (p != 0)
              {
              contoured = true;
              goto contour;
              }
            }
          }
      contour:
        if (contoured)
          {
          for (i[xi] = 0; i[xi] != xdim; i[xi]++)
            {
            for (i[yi] = 0; i[yi] != ydim; i[yi]++)
              {
              const TPixel p = m_Image->GetPixel (i);

              if (p != 0)
                {
                idxLower[xi] = min (idxLower[xi], i[xi]);
                idxUpper[xi] = max (idxUpper[xi], i[xi]);
                idxLower[yi] = min (idxLower[yi], i[yi]);
                idxUpper[yi] = max (idxUpper[yi], i[yi]);
                img->SetPixel (i, IN_INTENSITY);
                }
              else
                {
                img->SetPixel (i, OUT_INTENSITY);
                }
              }
            }
          }
        else
          {
          for (i[xi] = 0; i[xi] != xdim; i[xi]++)
            {
            for (i[yi] = 0; i[yi] != ydim; i[yi]++)
              {
              img->SetPixel (i, DONT_CARE_INTENSITY);
              }
            }
          }
        }
      // We are done with this image, free up some memory
      m_Image = NULL;
    }
    cout << ". done" << endl;
    cout << idxLower[0] << " " << idxLower[1] << " " << idxLower[2] << endl;
    cout << idxUpper[0] << " " << idxUpper[1] << " " << idxUpper[2] << endl;

    //
    // Now initialize the snake region
    //
    cout << "Initializing the level set region ... ";
    typename FloatImageType::Pointer imgLevelSet = FloatImageType::New ();

    imgLevelSet->SetRegions (size);
    imgLevelSet->SetSpacing (img->GetSpacing ());
    imgLevelSet->Allocate ();
    imgLevelSet->FillBuffer (+1.0);

    {
      FloatImageType::SizeType initRegionSize;
      initRegionSize[0] = 1 + idxUpper[0] - idxLower[0];
      initRegionSize[1] = 1 + idxUpper[1] - idxLower[1];
      initRegionSize[2] = 1 + idxUpper[2] - idxLower[2];

      FloatImageType::RegionType initRegion (idxLower, initRegionSize);
      ImageRegionIterator < FloatImageType > itLevelSet (imgLevelSet,
                                                         initRegion);
      ImageRegionIterator < FloatImageType > itImg (img,
                                                    initRegion);
      // Fill in the "initial snake"
      while (!itLevelSet.IsAtEnd ())
        {
        itLevelSet.Value () = ( itImg.Value() >= 0.0 ) ? -1.0 : 0.0;
        ++itLevelSet;
		++itImg;
        }
    }
    cout << "done" << endl;

    //
    // Initialize the snake driver
    //
    if( debuggingLevel > 0 )
      {
      saveImage( (FloatImageType*)imgLevelSet, "level-set.raw" );
      saveImage( (FloatImageType*)img, "speed-snapshot.raw" );
      }
    cout << "Initializing the level set driver ... ";
    SnakeParameters sp = SnakeParameters::GetDefaultInOutParameters ();

    sp.SetPropagationWeight (params["PropagationWeight"][0.5]);
    sp.SetCurvatureWeight (params["CurvatureWeight"][0.5]);

    SNAPLevelSetDriver3d *m_LevelSetDriver =
      new SNAPLevelSetDriver3d (imgLevelSet, // snake initialization
                                img,         // speed image
                                sp           // snake parameters
      );

    cout << "done" << endl;

    //
    // Now run the snake
    //
    {
      const unsigned int nIterations =
        (unsigned int) params["SnakeIterations"][10];
      for (int i = 0; i != nIterations; ++i)
        {
        m_LevelSetDriver->Run (1);
        cout << "\rRunning : Iteration " << m_LevelSetDriver->
            GetElapsedIterations () << "/" << nIterations << " ... ";
        m_LevelSetDriver->GetCurrentState()->Update();
        }
      // We are done with the speed image, free up some memory
      img = NULL;
    }
    cout << "done" << endl;
    if( debuggingLevel > 0 )
      {
      saveImage( (FloatImageType*)m_LevelSetDriver->GetCurrentState(), "snake-final.raw" );
      }

    //
    // Get the result.
    //
    cout << "Extracting the level set image ... ";
    LabelImageType::Pointer imgLabel = LabelImageType::New ();
    imgLabel->SetRegions (size);
    imgLabel->SetSpacing (imgLevelSet->GetSpacing ());
    imgLabel->Allocate ();

    {
      ImageRegionConstIterator < FloatImageType >
        itLevelSet (m_LevelSetDriver->GetCurrentState (),
                    m_LevelSetDriver->GetCurrentState ()->
                    GetLargestPossibleRegion ());
      ImageRegionIterator <LabelImageType> itSeg (imgLabel,
                                                  imgLabel->
                                                  GetLargestPossibleRegion());
      while (!itLevelSet.IsAtEnd ())
        {
          // Interior is negative.
          itSeg.Value () =
            (itLevelSet.Value () < 0.0) ? LabelType (1) : LabelType (0);
          ++itSeg;
          ++itLevelSet;
        }
      // Free up everything we are done with
      delete m_LevelSetDriver;

      // We are done with the level set image, free up some memory.
      imgLevelSet = NULL;
    }
    cout << "done" << endl;
    saveImage ((LabelImageType *) imgLabel,
               params["SegmentationFile"][params["GreyImage"]["no such file"]
                                          + std::string (".out")].c_str ());

    //
    // Extract the mesh
    //
    cout << "Extracting mesh ... ";
    {
      std::string fname =
        params["SegmentationFile"][params["GreyImage"]["no such file"] +
                                   std::string (".byu")].c_str ();
      // Create a pipeline for mesh generation
      IRISMeshPipeline *meshPipeline = new IRISMeshPipeline ();

      meshPipeline->SetImage (imgLabel);

      MeshOptions mo;

      mo.SetUseGaussianSmoothing(params["byuUseGaussianSmoothing"][true]);
      mo.SetGaussianStandardDeviation(params["byuGaussianStandardDeviation"]
                                      [0.8]);
      mo.SetGaussianError(params["byuGaussianError"][0.03]);
      mo.SetUseDecimation(params["byuUseDecimation"][true]);
      mo.SetDecimateTargetReduction(params["byuDecimateTargetReduction"][0.1]);
      mo.SetDecimateInitialError(params["byuDecimateInitialError"][0.002]);
      mo.SetDecimateAspectRatio(params["byuDecimateAspectRatio"][20.0]);
      mo.SetDecimateFeatureAngle(params["byuDecimateFeatureAngle"][45.0]);
      mo.SetDecimateErrorIncrement(params["byuDecimateErrorIncrement"][0.002]);
      mo.SetDecimateMaximumIterations(params["byuDecimateMaximumIterations"]
                                      [20]);
      mo.SetDecimatePreserveTopology(params["byuDecimatePreserveTopology"]
                                     [true]);
      mo.SetUseMeshSmoothing(params["byuUseMeshSmoothing"][true]);
      mo.SetMeshSmoothingRelaxationFactor(params
                                          ["byuMeshSmoothingRelaxationFactor"]
                                          [0.01]);
      mo.SetMeshSmoothingIterations(params["byuMeshSmoothingIterations"][5]);
      mo.SetMeshSmoothingConvergence(params["byuMeshSmoothingConvergence"]
                                     [0.0]);
      mo.SetMeshSmoothingFeatureEdgeSmoothing(params
                                              ["byuMeshSmoothingFeatureEdgeSmoothing"]
                                              [true]);
      mo.SetMeshSmoothingBoundarySmoothing(params
                                           ["byuMeshSmoothingBoundarySmoothing"]
                                           [true]);
      // Pass the settings on to the pipeline
      meshPipeline->SetMeshOptions (mo);

      // Run the first step in this pipeline
      meshPipeline->ComputeBoundingBoxes ();

      vtkPolyData *mesh = vtkPolyData::New ();

      // 1 is the label we are using for the segmentated result.
      meshPipeline->ComputeMesh (1, mesh);

      vtkTriangleFilter *triangleFilter = vtkTriangleFilter::New ();

      triangleFilter->SetInput (mesh);
      triangleFilter->PassLinesOn ();
      triangleFilter->PassVertsOn ();

      vtkBYUWriter *byu = vtkBYUWriter::New ();

      byu->SetInput (triangleFilter->GetOutput ());
      byu->SetGeometryFileName (fname.c_str ());
      byu->WriteTextureOff ();
      byu->WriteScalarOff ();
      byu->WriteDisplacementOff ();
      byu->Update ();

      // Deallocate the filter
      delete meshPipeline;
    }
    cout << "done" << endl;
  }

private:
  /**
   * The currently loaded image.
   */
  ImagePointer m_Image;

  /**
   * The image IO mechanism
   */
  ImageIOPointer m_ImageIO;

  /**
   * The set of options for the run obtained from the command fle
   */
  Registry params;

  /**
   * The debugging level
   * At level 0, no intermediate images are saved.
   * At level 1, intermediate images are saved.
   */
   int debuggingLevel;
};


/**
 * Parse the command line and pass on the command file name to the 
 */
int
main (int argc, char **argv)
{
  // Parse command line parameters
  CommandLineArgumentParser parser;

  // To allow for automation.
  parser.AddOption ("--parameters", 1);
  parser.AddSynonim ("--parameters", "-p");
  parser.AddSynonim ("--parameters", "--param");

  parser.AddOption ("--greyImage", 1);

  CommandLineArgumentParseResult parseResult;

  if (!parser.TryParseCommandLine (argc, argv, parseResult))
    {
    // Print usage info and exit
    cerr << endl
	     << "Contour Interpolation with Geodesic Snakes Command Line Usage:" << endl
         << "   GeoInterp [options]" << endl;

    cerr << "Options:" << endl;
    cerr << "   --parameters, -p FILE: "
         << "Load segmentation parameters from file FILE" << endl
	     << "   --greyImage FILE     : "
		 << "The image filename (overrides the parameter file)" << endl;
    return -1;
    }

  // Initialize the operating system interface
  // SystemInterface system;  

  if (parseResult.IsOptionPresent ("--parameters"))
    {
    // Get the filename
    const char *fname     = parseResult.GetOptionParameter ("--parameters");
	const char *greyImage = NULL;

    if (parseResult.IsOptionPresent ("--greyImage"))
	  {
	  greyImage           = parseResult.GetOptionParameter ("--greyImage");
      }

    Automate < GreyType > automate;
    automate.RunSegmentation (fname, greyImage);
    }
  return 0;
}

// vi: set ts=2 sw=2 si: