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

  This file is part of the elastix software.

  Copyright (c) University Medical Center Utrecht. All rights reserved.
  See src/CopyrightElastix.txt or http://elastix.isi.uu.nl/legal.php for
  details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE. See the above copyright notices for more information.

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

#ifndef __elxFullSearchOptimizer_hxx
#define __elxFullSearchOptimizer_hxx

#include "elxFullSearchOptimizer.h"
#include <iomanip>
#include <sstream>
#include <string>
#include "vnl/vnl_math.h"

namespace elastix
{
using namespace itk;


/**
 * ********************* Constructor ****************************
 */

template <class TElastix>
FullSearch<TElastix>
::FullSearch()
{
  this->m_OptimizationSurface = 0;

} // end Constructor


/**
 * ***************** BeforeRegistration ***********************
 */

template <class TElastix>
void
FullSearch<TElastix>
::BeforeRegistration( void )
{
  /** Add the target cells "ItNr" and "Metric" to xout["iteration"]. */
  xl::xout["iteration"].AddTargetCell("2:Metric");

  /** Format the metric as floats. */
  xl::xout["iteration"]["2:Metric"] << std::showpoint << std::fixed;

} // end BeforeRegistration


/**
 * ***************** BeforeEachResolution ***********************
 */

template <class TElastix>
void
FullSearch<TElastix>
::BeforeEachResolution( void )
{
  /** Get the current resolution level.*/
  unsigned int level = static_cast<unsigned int>(
    this->m_Registration->GetAsITKBaseType()->GetCurrentLevel() );

  /** Read FullSearchRange from the parameter file. */

  /** declare variables */
  std::string name("");
  unsigned int param_nr = 0;
  double minimum = 0;
  double maximum = 0;
  double stepsize = 0;
  const std::string prefix = "FS";
  unsigned int entry_nr = 0;
  unsigned int nrOfSearchSpaceDimensions = 0;
  bool found = true;
  bool realGood = true;
  std::ostringstream makeString( "" );

  /** Create fullFieldName, which is "FullSearchSpace0" at level 0. */
  makeString << "FullSearchSpace" << level;
  std::string fullFieldName = makeString.str();

  /** Loop as long as search ranges are defined. */
  while ( found )
  {
    /** Try to read (silently) from the parameter file. */
    /** \todo check earlier, in BeforeAll, if the searchspace has been defined. */

    if ( realGood && found )
    {
      found = this->GetConfiguration()->ReadParameter(
        name, fullFieldName.c_str(), entry_nr, false );
      realGood = this->CheckSearchSpaceRangeDefinition(
        fullFieldName.c_str(), found, entry_nr );
      entry_nr++;
    }
    if ( realGood && found )
    {
      found = this->GetConfiguration()->ReadParameter(
        param_nr, fullFieldName.c_str(), entry_nr, false );
      realGood = this->CheckSearchSpaceRangeDefinition(
        fullFieldName.c_str(), found, entry_nr );
      entry_nr++;
    }
    if ( realGood && found )
    {
      found = this->GetConfiguration()->ReadParameter(
        minimum, fullFieldName.c_str(), entry_nr, false );
      realGood = this->CheckSearchSpaceRangeDefinition(
        fullFieldName.c_str(), found, entry_nr );
      entry_nr++;
    }
    if ( realGood && found )
    {
      found = this->GetConfiguration()->ReadParameter(
        maximum, fullFieldName.c_str(), entry_nr, false );
      realGood = this->CheckSearchSpaceRangeDefinition(
        fullFieldName.c_str(), found, entry_nr );
      entry_nr++;
    }
    if ( realGood && found )
    {
      found = this->GetConfiguration()->ReadParameter(
        stepsize, fullFieldName.c_str(), entry_nr, false );
      realGood = this->CheckSearchSpaceRangeDefinition(
        fullFieldName.c_str(), found, entry_nr );
      entry_nr++;
    }

    /** Setup this search range. */
    if ( realGood && found )
    {
      /** Setup the Superclass */
      this->AddSearchDimension( param_nr, minimum, maximum, stepsize );

      /** Create name of this dimension. */
      makeString.str( "" );
      makeString
        << prefix
        << ( (entry_nr / 5) - 1 )
        << ":"
        << name
        << ":"
        << param_nr;

      /** Store the name and create a column in xout["iteration"]. */
      this->m_SearchSpaceDimensionNames[param_nr] = makeString.str();
      xl::xout["iteration"].AddTargetCell( makeString.str().c_str() );

      /** Format this xout iteration column as float. */
      xl::xout["iteration"][ makeString.str().c_str() ]
        << std::showpoint << std::fixed;
    }
  } // end while

  if ( realGood )
  {
    /** The number of dimensions. */
    nrOfSearchSpaceDimensions = this->GetNumberOfSearchSpaceDimensions();

    /** Create the image that will store the results of the full search. */
    this->m_OptimizationSurface
      = NDImageType::NewNDImage( nrOfSearchSpaceDimensions );
    this->m_OptimizationSurface->CreateNewImage();
    /** \todo don't do this if more than max allowable dimensions. */

    /** Set the correct size and allocate memory. */
    this->m_OptimizationSurface->SetRegions(
      this->GetSearchSpaceSize()  );
    this->m_OptimizationSurface->Allocate();
    /** \todo try/catch block around Allocate? */

    /** Set the name of this image on disk. */
    std::string resultImageFormat = "mhd";
    this->m_Configuration->ReadParameter(
      resultImageFormat, "ResultImageFormat", 0, false );
    makeString.str("");
    makeString
      << this->GetConfiguration()->GetCommandLineArgument( "-out" )
      << "OptimizationSurface."
      << this->GetConfiguration()->GetElastixLevel()
      << ".R" << level
      << "." << resultImageFormat;
    this->m_OptimizationSurface->SetOutputFileName( makeString.str().c_str() );

    elxout
      << "Total number of iterations needed in this resolution: "
      << this->GetNumberOfIterations()
      << "." << std::endl;

  }
  else
  {
    itkExceptionMacro( << "ERROR: elastix found an error in the search "
      << "space definition, and is quiting." );
  }

} // end BeforeEachResolution()


/**
 * ***************** AfterEachIteration *************************
 */

template <class TElastix>
void
FullSearch<TElastix>
::AfterEachIteration( void )
{
  /** Print some information. */
  xl::xout["iteration"]["2:Metric"] << this->GetValue();

  this->m_OptimizationSurface->SetPixel(
    this->GetCurrentIndexInSearchSpace(), this->GetValue() );

  SearchSpacePointType currentPoint = this->GetCurrentPointInSearchSpace();
  unsigned int nrOfSSDims = currentPoint.GetSize();
  NameIteratorType name_it = this->m_SearchSpaceDimensionNames.begin();

  for ( unsigned int dim = 0; dim < nrOfSSDims; dim++ )
  {
    xl::xout["iteration"][ name_it->second.c_str() ] << currentPoint[ dim ];
    name_it++;
  }

} // end AfterEachIteration()


/**
 * ***************** AfterEachResolution *************************
 */

template <class TElastix>
void
FullSearch<TElastix>
::AfterEachResolution( void )
{
  //typedef enum {FullRangeSearched,  MetricError } StopConditionType;
  std::string stopcondition;

  switch( this->GetStopCondition() )
  {
  case FullRangeSearched :
    stopcondition = "The full range has been searched";
    break;

  case MetricError :
    stopcondition = "Error in metric";
    break;

  default:
    stopcondition = "Unknown";
    break;
  }

  /** Print the stopping condition */
  elxout << "Stopping condition: " << stopcondition << "." << std::endl;

  /** Write the optimization surface to disk */
  try
  {
    this->m_OptimizationSurface->Write();
    elxout
      << "\nThe scanned optimization surface is saved as: "
      << this->m_OptimizationSurface->GetOutputFileName()
      << std::endl;
  }
  catch ( ExceptionObject& err )
  {
    xl::xout["error"]
      << "ERROR: Saving "
      << this->m_OptimizationSurface->GetOutputFileName()
      << " failed."
      << std::endl;
    xl::xout["error"] << err << std::endl;
    // do not throw an error, since we would like to go on.
  }

  /** Print the best metric value */
  elxout
    << std::endl
    << "Best metric value in this resolution = "
    << this->GetBestValue()
    << std::endl;

  /** Print the best index and point */
  SearchSpaceIndexType bestIndex = this->GetBestIndexInSearchSpace();
  SearchSpacePointType bestPoint = this->GetBestPointInSearchSpace();
  unsigned int nrOfSSDims = bestIndex.GetSize();

  elxout << "Index of the point in the optimization surface image that has "
    << "the best metric value: [ ";
  for ( unsigned int dim = 0; dim < nrOfSSDims; dim++ )
  {
    elxout << bestIndex[ dim ] << " ";
  }
  elxout << "]" << std::endl;

  elxout << "The corresponding parameter values: [ " ;
  for ( unsigned int dim = 0; dim < nrOfSSDims; dim++ )
  {
    elxout << bestPoint[ dim ] << " ";
  }
  elxout << "]\n" << std::endl;

  /** Remove the columns from xout["iteration"]. */
  NameIteratorType name_it = this->m_SearchSpaceDimensionNames.begin();
  for ( unsigned int dim = 0; dim < nrOfSSDims; dim++ )
  {
    xl::xout["iteration"].RemoveTargetCell( name_it->second.c_str() );
    name_it++;
  }

  /** Clear the dimension names of the previous resolution's search space. */
  this->m_SearchSpaceDimensionNames.clear();

  /** Clear the full search ranges */
  this->SetSearchSpace( 0 );

} // end AfterEachResolution()


/**
 * ******************* AfterRegistration ************************
 */
template <class TElastix>
void
FullSearch<TElastix>
::AfterRegistration( void )
{
  /** Print the best metric value. */
  double bestValue = this->GetBestValue();
  elxout << std::endl << "Final metric value  = " << bestValue << std::endl;

} // end AfterRegistration()


/**
 * ************ CheckSearchSpaceRangeDefinition *****************
 */

template <class TElastix>
bool
FullSearch<TElastix>
::CheckSearchSpaceRangeDefinition( const std::string & fullFieldName,
  const bool found, const unsigned int entry_nr ) const
{
  /** Complain if not at least one search space dimension has been found,
   * or if a search dimension is not fully specified.
   */
  if ( !found && ( entry_nr == 0 || ( entry_nr % 5 != 0 ) ) )
  {
    xl::xout["error"]
      << "ERROR:\nNo (valid) range specified for the full search optimizer!\n"
      << "Please define the field ("
      << fullFieldName
      << " \"name\" parameter_nr min max stepsize) correctly in the "
      << "parameter file"
      << std::endl;
    return false;
  }

  return true;

} // end CheckSearchSpaceRangeDefinition()


} // end namespace elastix

#endif // end #ifndef __elxFullSearchOptimizer_hxx

