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

  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 __itkRSGDEachParameterApartBaseOptimizer_cxx
#define __itkRSGDEachParameterApartBaseOptimizer_cxx

#include "itkRSGDEachParameterApartBaseOptimizer.h"
#include "itkCommand.h"
#include "itkEventObject.h"
#include "vnl/vnl_math.h"

namespace itk
{

/**
 * Constructor
 */
RSGDEachParameterApartBaseOptimizer
::RSGDEachParameterApartBaseOptimizer()
{

  itkDebugMacro("Constructor");

  m_MaximumStepLength = 1.0;
  m_MinimumStepLength = 1e-3;
  m_GradientMagnitudeTolerance = 1e-4;
  m_NumberOfIterations = 100;
  m_CurrentIteration   =   0;
  m_Value = 0;
  m_Maximize = false;
  m_CostFunction = 0;

  m_CurrentStepLengths.Fill( 0.0f );
  m_CurrentStepLength = 0;

  m_StopCondition = MaximumNumberOfIterations;
  m_Gradient.Fill( 0.0f );
  m_PreviousGradient.Fill( 0.0f );

  m_GradientMagnitude = 0.0;

}



/**
 * Start the optimization
 */
void
RSGDEachParameterApartBaseOptimizer
::StartOptimization( void )
{

  itkDebugMacro("StartOptimization");

  m_CurrentIteration          = 0;

  const unsigned int spaceDimension =
    m_CostFunction->GetNumberOfParameters();

  m_Gradient = DerivativeType( spaceDimension );
  m_PreviousGradient = DerivativeType( spaceDimension );
  m_CurrentStepLengths = DerivativeType( spaceDimension );
  m_Gradient.Fill( 0.0f );
  m_PreviousGradient.Fill( 0.0f );
  m_CurrentStepLengths.Fill( m_MaximumStepLength );
  m_CurrentStepLength = m_MaximumStepLength;

  this->SetCurrentPosition( GetInitialPosition() );
  this->ResumeOptimization();

}





/**
 * Resume the optimization
 */
void
RSGDEachParameterApartBaseOptimizer
::ResumeOptimization( void )
{

  itkDebugMacro("ResumeOptimization");

  m_Stop = false;

  this->InvokeEvent( StartEvent() );

  while( !m_Stop )
    {

    /** inefficient:
    ParametersType currentPosition = this->GetCurrentPosition();
    m_Value = m_CostFunction->GetValue( currentPosition );

    if( m_Stop )
      {
      break;
      }

    m_PreviousGradient = m_Gradient;
    m_CostFunction->GetDerivative( currentPosition, m_Gradient );
    */

    /** faster:*/

    m_PreviousGradient = m_Gradient;
    try
    {
      m_CostFunction->GetValueAndDerivative(
        this->GetCurrentPosition(), m_Value, m_Gradient );
    }
    catch( ExceptionObject& err)
    {
      m_StopCondition = MetricError;
      this->StopOptimization();
      throw err;
    }


    if( m_Stop )
      {
      break;
      }

    this->AdvanceOneStep();

    m_CurrentIteration++;

    if( m_CurrentIteration == m_NumberOfIterations )
      {
      m_StopCondition = MaximumNumberOfIterations;
      this->StopOptimization();
      break;
      }

    }


}





/**
 * Stop optimization
 */
void
RSGDEachParameterApartBaseOptimizer
::StopOptimization( void )
{

  itkDebugMacro("StopOptimization");

  m_Stop = true;
  this->InvokeEvent( EndEvent() );
}




/**
 * Advance one Step following the gradient direction
 */
void
RSGDEachParameterApartBaseOptimizer
::AdvanceOneStep( void )
{

  itkDebugMacro("AdvanceOneStep");

  const unsigned int  spaceDimension =
    m_CostFunction->GetNumberOfParameters();

  DerivativeType transformedGradient( spaceDimension );
  DerivativeType previousTransformedGradient( spaceDimension );
  ScalesType     scales = this->GetScales();

  // Make sure the scales have been set properly
  if (scales.size() != spaceDimension)
    {
    itkExceptionMacro(<< "The size of Scales is "
                      << scales.size()
                      << ", but the NumberOfParameters for the CostFunction is "
                      << spaceDimension
                      << ".");
    }

  for(unsigned int i = 0;  i < spaceDimension; i++)
    {
    transformedGradient[i]  = m_Gradient[i] / scales[i];
    previousTransformedGradient[i] =
      m_PreviousGradient[i] / scales[i];
    }

  double magnitudeSquare = 0;
  for(unsigned int dim=0; dim<spaceDimension; dim++)
    {
    const double weighted = transformedGradient[dim];
    magnitudeSquare += weighted * weighted;
    }

  m_GradientMagnitude = vcl_sqrt( magnitudeSquare );

  if( m_GradientMagnitude < m_GradientMagnitudeTolerance )
    {
    m_StopCondition = GradientMagnitudeTolerance;
    StopOptimization();
    return;
    }

  double sumOfCurrentStepLengths = 0.0;
  double biggestCurrentStepLength = 0.0;
  for(unsigned int i=0; i<spaceDimension; i++)
  {
    const bool signChange =
      ( transformedGradient[i] * previousTransformedGradient[i] ) < 0 ;

    if (signChange)
    {
      m_CurrentStepLengths[i] /=2.0;
    }

    const double currentStepLengths_i = m_CurrentStepLengths[i];

    sumOfCurrentStepLengths += currentStepLengths_i;
    if ( currentStepLengths_i > biggestCurrentStepLength )
    {
      biggestCurrentStepLength = currentStepLengths_i;
    }
  } //end for

  /** The average current step length: */
  m_CurrentStepLength = sumOfCurrentStepLengths / spaceDimension;

  /** if all current step lengths are smaller than the
   * MinimumStepLength stop the optimization
   */
  if( biggestCurrentStepLength < m_MinimumStepLength )
    {
    m_StopCondition = StepTooSmall;
    StopOptimization();
    return;
    }

  double direction;
  if( this->m_Maximize )
    {
    direction = 1.0;
    }
  else
    {
    direction = -1.0;
    }

  DerivativeType factor = DerivativeType(spaceDimension);

  for(unsigned int i=0; i<spaceDimension; i++)
  {
    factor[i] = direction * m_CurrentStepLengths[i] / m_GradientMagnitude;
  }

  // This method StepAlongGradient() will
  // be overloaded in non-vector spaces
  this->StepAlongGradient( factor, transformedGradient );



  this->InvokeEvent( IterationEvent() );

}

void
RSGDEachParameterApartBaseOptimizer
::PrintSelf( std::ostream& os, Indent indent ) const
{
  Superclass::PrintSelf(os,indent);
  os << indent << "MaximumStepLength: "
     << m_MaximumStepLength << std::endl;
  os << indent << "MinimumStepLength: "
     << m_MinimumStepLength << std::endl;
  os << indent << "GradientMagnitudeTolerance: "
     << m_GradientMagnitudeTolerance << std::endl;
  os << indent << "NumberOfIterations: "
     << m_NumberOfIterations << std::endl;
  os << indent << "CurrentIteration: "
     << m_CurrentIteration   << std::endl;
  os << indent << "Value: "
     << m_Value << std::endl;
  os << indent << "Maximize: "
     << m_Maximize << std::endl;
  if (m_CostFunction)
    {
    os << indent << "CostFunction: "
       << &m_CostFunction << std::endl;
    }
  else
    {
    os << indent << "CostFunction: "
       << "(None)" << std::endl;
    }
  os << indent << "CurrentStepLength: "
     << m_CurrentStepLength << std::endl;
  os << indent << "StopCondition: "
     << m_StopCondition << std::endl;
  os << indent << "Gradient: "
     << m_Gradient << std::endl;
}
} // end namespace itk

#endif
