// -------------------------------------------------------------------------
// @author Leonardo Florez-Valencia (florez-l@javeriana.edu.co)
// -------------------------------------------------------------------------

#ifndef __CPEXTENSIONS__ALGORITHMS__ITERATIVEGAUSSIANMODELESTIMATOR__HXX__
#define __CPEXTENSIONS__ALGORITHMS__ITERATIVEGAUSSIANMODELESTIMATOR__HXX__

#include <cmath>
#include <cstdarg>
#include <fstream>
#include <vnl/vnl_math.h>
#include <vnl/vnl_inverse.h>

// -------------------------------------------------------------------------
template< class S, unsigned int D >
void cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
SetNumberOfSamples( unsigned long n )
{
  this->m_N = S( n );
  this->m_Updated = false;
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
void cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
SetMu( const TMatrix& m )
{
  if( m.rows( ) == D && m.columns( ) == 1 )
  {
    this->m_M = m;
    this->m_Updated = false;
    this->Modified( );
  }
  else
  {
    itkExceptionMacro(
      << "Input Mu matrix is not a " << D << "x1 matrix"
      );

  } // fi
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
void cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
SetOmega( const TMatrix& O )
{
  if( O.rows( ) == D && O.columns( ) == D )
  {
    this->m_O = O;
    this->m_Updated = false;
    this->Modified( );
  }
  else
  {
    itkExceptionMacro(
      << "Input Omega matrix is not a " << D << "x" << D << " matrix"
      );

  } // fi
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
bool cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
SaveModelToFile( const std::string& filename ) const
{
  std::ofstream out( filename.c_str( ) );
  if( out )
  {
    out << this->m_N << std::endl;
    out << this->m_M << std::endl;
    out << this->m_O << std::endl;
    out.close( );
    return( true );
  }
  else
    return( false );
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
bool cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
LoadModelFromFile( const std::string& filename )
{
  std::ifstream in( filename.c_str( ) );
  if( in )
  {
    this->Clear( );
    in >> this->m_N;
    for( unsigned int i = 0; i < D; ++i )
      in >> this->m_M[ i ][ 0 ];
    for( unsigned int j = 0; j < D; ++j )
      for( unsigned int i = 0; i < D; ++i )
        in >> this->m_O[ i ][ j ];
    in.close( );
    return( true );
  }
  else
    return( false );
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
template< class V >
S cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
Probability( const V& sample ) const
{
  if( !this->m_Updated )
    this->_UpdateModel( );

  TMatrix c( D, 1 );
  for( unsigned int d = 0; d < D; ++d )
    c[ d ][ 0 ] = S( sample[ d ] ) - this->m_M[ d ][ 0 ];
  if( S( 0 ) < this->m_Norm )
  {
    // Covariance is NOT null
    double v = double( ( c.transpose( ) * ( this->m_Inv * c ) )[ 0 ][ 0 ] );
    v /= double( 2 );

    return( S( std::exp( -v ) ) );
  }
  else
  {
    // Covariance is null
    S n = S( 0 );
    for( unsigned int d = 0; d < D; ++d )
      n += c[ d ][ 0 ] * c[ d ][ 0 ];
    bool equal = ( double( n ) < double( 1e-10 ) );
    return( ( equal )? S( 1 ): S( 0 ) );

  } // fi
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
S cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
Probability( const S& s_x, const S& s_y, ... ) const
{
  static S sample[ D ];

  std::va_list args_lst;
  va_start( args_lst, s_y );
  sample[ 0 ] = s_x;
  sample[ 1 ] = s_y;
  for( unsigned int d = 2; d < D; ++d )
    sample[ d ] = S( va_arg( args_lst, double ) );
  va_end( args_lst );
  return( this->Probability( sample ) );
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
template< class V, class M >
void cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
GetModel( V& m, M& E ) const
{
  if( !this->m_Updated )
    this->_UpdateModel( );
  for( unsigned int i = 0; i < D; ++i )
  {
    m[ i ] = double( this->m_M[ i ][ 0 ] );
    for( unsigned int j = 0; j < D; ++j )
      E[ i ][ j ] = double( this->m_Cov[ i ][ j ] );

  } // rof
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
void cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
Clear( )
{
  this->m_N = S( 0 );
  this->m_M.set_size( D, 1 );
  this->m_O.set_size( D, D );
  this->m_M.fill( S( 0 ) );
  this->m_O.fill( S( 0 ) );
  this->m_Updated = false;
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
template< class V >
void cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
AddSample( const V& sample )
{
  TMatrix s( D, 1 );
  for( unsigned int d = 0; d < D; ++d )
    s[ d ][ 0 ] = S( sample[ d ] );

  // Update mean
  S coeff = this->m_N;
  this->m_N += S( 1 );
  coeff /= this->m_N;
  this->m_M = ( this->m_M * coeff ) + ( s / this->m_N );

  // Update omega operand
  if( this->m_N == 1 )
    this->m_O  = s * s.transpose( );
  else if( this->m_N == 2 )
    this->m_O += s * s.transpose( );
  else
  {
    S inv = S( 1 ) / ( this->m_N - S( 1 ) );
    this->m_O  = this->m_O * ( this->m_N - S( 2 ) ) * inv;
    this->m_O += ( s * s.transpose( ) ) * inv;

  } // fi

  this->m_Updated = false;
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
void cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
AddSample( const S& s_x, const S& s_y, ... )
{
  static S sample[ D ];

  std::va_list args_lst;
  va_start( args_lst, s_y );
  sample[ 0 ] = s_x;
  if( D > 1 )
  {
    sample[ 1 ] = s_y;
    for( unsigned int d = 2; d < D; ++d )
      sample[ d ] = S( va_arg( args_lst, double ) );
    va_end( args_lst );

  } // fi
  this->AddSample( sample );
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
IterativeGaussianModelEstimator( )
  : Superclass( )
{
  this->Clear( );
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
~IterativeGaussianModelEstimator( )
{
}

// -------------------------------------------------------------------------
template< class S, unsigned int D >
void cpExtensions::Algorithms::IterativeGaussianModelEstimator< S, D >::
_UpdateModel( ) const
{
  static const double _2piD =
    std::pow( double( 2 ) * double( vnl_math::pi ), D );

  // Update covariance
  this->m_Cov =
    this->m_O -
    (
      ( this->m_M * this->m_M.transpose( ) ) *
      ( this->m_N / ( this->m_N - S( 1 ) ) )
      );

  // Compute inverse and determinant
  S det = vnl_determinant( this->m_Cov );
  if( !( det > S( 0 ) ) )
  {
    this->m_Inv.set_size( D, D );
    this->m_Inv.fill( S( 0 ) );
    this->m_Norm = S( 0 );
  }
  else
  {
    this->m_Inv = vnl_inverse( this->m_Cov );
    this->m_Norm = S( 1 ) / S( std::sqrt( _2piD * double( det ) ) );

  } // fi

  // Object now is updated
  this->m_Updated = true;
}

#endif // __CPEXTENSIONS__ALGORITHMS__ITERATIVEGAUSSIANMODELESTIMATOR__HXX__

// eof - $RCSfile$