// =========================================================================
// @author Leonardo Florez Valencia
// @email florez-l@javeriana.edu.co
// =========================================================================
#ifndef __fpa__Common__OriginalRandomWalker__hxx__
#define __fpa__Common__OriginalRandomWalker__hxx__

#include <cmath>
#include <map>
#include <itkImageRegionConstIteratorWithIndex.h>
#include <Eigen/Sparse>

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalar >
void fpa::Common::OriginalRandomWalker< _TImage, _TLabels, _TScalar >::
AddSeed( const TIndex& seed, const TLabel& label )
{
  this->m_Seeds.push_back( seed );
  this->m_Labels.push_back( label );
  this->Modified( );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalar >
fpa::Common::OriginalRandomWalker< _TImage, _TLabels, _TScalar >::
OriginalRandomWalker( )
  : Superclass( ),
    m_Beta( TScalar( 90 ) ),
    m_Epsilon( TScalar( 1e-5 ) ),
    m_NormalizeWeights( true )
{
  fpaFilterOptionalInputConfigureMacro( InputLabels, TLabels );
  fpaFilterOutputConfigureMacro( OutputProbabilities, TScalarImage );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalar >
fpa::Common::OriginalRandomWalker< _TImage, _TLabels, _TScalar >::
~OriginalRandomWalker( )
{
}

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalar >
void fpa::Common::OriginalRandomWalker< _TImage, _TLabels, _TScalar >::
GenerateData( )
{
  typedef Eigen::SparseMatrix< TScalar >           _TSparseMatrix;
  typedef Eigen::SimplicialLDLT< _TSparseMatrix >  _TSparseSolver;
  typedef Eigen::Triplet< TScalar >                _TTriplet;
  typedef std::vector< _TTriplet >                 _TTriplets;
  typedef std::map< unsigned long, unsigned long > _TMap;
  typedef _TMap::value_type                        _TMapValue;

  // Some input values
  const TImage* in = this->GetInput( );
  TRegion region = in->GetRequestedRegion( );
  unsigned long N = region.GetNumberOfPixels( );

  // Prepare seeds for linear algebra
  _TMap seeds;
  for( unsigned int i = 0; i < this->m_Seeds.size( ); ++i )
    seeds.insert(
      _TMapValue(
        Self::_1D( this->m_Seeds[ i ], region ),
        this->m_Labels[ i ]
        )
      );

  // Use input labels
  /* TODO
     const TLabels* in_labels = this->GetInputLabels( );
     if( in_labels != NULL )
     {
     itk::ImageRegionConstIteratorWithIndex< TLabels > lIt( in_labels, region );
     for( lIt.GoToBegin( ); !lIt.IsAtEnd( ); ++lIt )
     if( lIt.Get( ) != 0 )
     seeds.insert(
     _TMapValue( Self::_1D( lIt.GetIndex( ), region ), lIt.Get( ) )
     );

     } // fi
  */

  // Prepare label tables
  _TMap labels, inv_labels, seeds_indexes;
  _TMap::const_iterator hIt = seeds.begin( );
  for( ; hIt != seeds.end( ); ++hIt )
  {
    seeds_indexes[ hIt->first ] = seeds_indexes.size( ) - 1;
    if( labels.find( hIt->second ) == labels.end( ) )
    {
      labels[ hIt->second ] = labels.size( ) - 1;
      inv_labels[ labels[ hIt->second ] ] = hIt->second;

    } // fi

  } // rof

  // Prepare matrix/image index conversion
  _TMap indexes;
  unsigned long o = 0;
  for( unsigned long n = 0; n < N; ++n )
  {
    if( seeds.find( n ) == seeds.end( ) )
      indexes.insert( _TMap::value_type( n, n - o ) );
    else
      o++;

  } // rof
  unsigned long nLabels = labels.size( );

  // Boundary matrix
  _TTriplets Bt;
  _TMap::const_iterator mIt = seeds.begin( );
  for( unsigned long i = 0; mIt != seeds.end( ); ++mIt, ++i )
    Bt.push_back( _TTriplet( i, labels[ mIt->second ], TScalar( 1 ) ) );
  _TSparseMatrix B( seeds.size( ), nLabels );
  B.setFromTriplets( Bt.begin( ), Bt.end( ) );
  B.makeCompressed( );

  // Laplacian matrix
  _TTriplets Lt;
  itk::ImageRegionConstIteratorWithIndex< TImage > it( in, region );
  TScalar maxV = TScalar( -1 );
  for( it.GoToBegin( ); !it.IsAtEnd( ); ++it )
  {
    TIndex idx = it.GetIndex( );
    TScalar vidx = TScalar( it.Get( ) );
    unsigned long iidx = Self::_1D( idx, region );

    // Neighbors
    TScalar s = TScalar( 0 );
    for( unsigned int d = 0; d < TImage::ImageDimension; ++d )
    {
      TIndex jdx;
      for( int l = -1; l <= 1; l += 2 )
      {
        jdx = idx;
        jdx[ d ] += l;
        if( region.IsInside( jdx ) )
        {
          TScalar vjdx = TScalar( in->GetPixel( jdx ) );
          unsigned long ijdx = Self::_1D( jdx, region );
          TScalar v = std::fabs( vidx - vjdx );
          Lt.push_back( _TTriplet( iidx, ijdx, v ) );
          if( maxV < v )
            maxV = v;

        } // fi

      } // rof

    } // rof

  } // rof
  std::vector< TScalar > diag( N, TScalar( 0 ) );
  typename _TTriplets::iterator tIt;
  TScalar betaV = -this->m_Beta;
  if( this->m_NormalizeWeights )
    betaV /= maxV;
  for( tIt = Lt.begin( ); tIt != Lt.end( ); ++tIt )
  {
    TScalar v = std::exp( betaV * tIt->value( ) );
    if( v < this->m_Epsilon )
      v = this->m_Epsilon;
    *tIt = _TTriplet( tIt->row( ), tIt->col( ), -v );
    diag[ tIt->col( ) ] += v;

  } // rof
  for( unsigned long i = 0; i < diag.size( ); ++i )
    Lt.push_back( _TTriplet( i, i, diag[ i ] ) );

  // Compute R and A
  _TTriplets Rt, At;
  for( tIt = Lt.begin( ); tIt != Lt.end( ); ++tIt )
  {
    _TMap::const_iterator cIt, rIt;
    cIt = seeds.find( tIt->col( ) );
    rIt = seeds.find( tIt->row( ) );
    if( cIt != seeds.end( ) )
    {
      if( rIt == seeds.end( ) )
      {
        _TMap::const_iterator iIt = indexes.find( tIt->row( ) );
        _TMap::const_iterator jIt = seeds_indexes.find( cIt->first );
        Rt.push_back( _TTriplet( iIt->second, jIt->second, -tIt->value( ) ) );

      } // fi
    }
    else
    {
      if( rIt == seeds.end( ) )
      {
        _TMap::const_iterator iIt = indexes.find( tIt->row( ) );
        _TMap::const_iterator jIt = indexes.find( tIt->col( ) );
        At.push_back( _TTriplet( iIt->second, jIt->second, tIt->value( ) ) );

      } // fi

    } // fi

  } // rof

  _TSparseMatrix R( N - seeds.size( ), seeds.size( ) );
  R.setFromTriplets( Rt.begin( ), Rt.end( ) );
  R.makeCompressed( );

  _TSparseMatrix A( N - seeds.size( ), N - seeds.size( ) );
  A.setFromTriplets( At.begin( ), At.end( ) );
  A.makeCompressed( );

  // Solve dirichlet problem
  _TSparseSolver solver;
  solver.compute( A );
  if( solver.info( ) != Eigen::Success )
  {
    std::cerr << "Error computing." << std::endl;
  } // fi
  _TSparseMatrix x = solver.solve( R * B );
  if( solver.info( ) != Eigen::Success )
  {
    std::cerr << "Error solving." << std::endl;
  } // fi

  // Fill outputs
  TLabels* out_labels = this->GetOutput( );
  out_labels->SetLargestPossibleRegion( in->GetLargestPossibleRegion( ) );
  out_labels->SetRequestedRegion( in->GetRequestedRegion( ) );
  out_labels->SetBufferedRegion( in->GetBufferedRegion( ) );
  out_labels->SetSpacing( in->GetSpacing( ) );
  out_labels->SetOrigin( in->GetOrigin( ) );
  out_labels->SetDirection( in->GetDirection( ) );
  out_labels->Allocate( );

  TScalarImage* out_probs = this->GetOutputProbabilities( );
  out_probs->SetLargestPossibleRegion( in->GetLargestPossibleRegion( ) );
  out_probs->SetRequestedRegion( in->GetRequestedRegion( ) );
  out_probs->SetBufferedRegion( in->GetBufferedRegion( ) );
  out_probs->SetSpacing( in->GetSpacing( ) );
  out_probs->SetOrigin( in->GetOrigin( ) );
  out_probs->SetDirection( in->GetDirection( ) );
  out_probs->Allocate( );

  mIt = seeds.begin( );
  o = 0;
  unsigned long j = 0;
  for( unsigned long j = 0; j < N; ++j )
  {
    TIndex idx;
    TLabel lbl;
    TScalar p;
    if( mIt->first != j )
    {
      idx = Self::_ND( j, region );
      TScalar maxP = x.coeffRef( j - o, 0 );
      unsigned long maxL = 0;
      for( unsigned long l = 1; l < nLabels; ++l )
      {
        TScalar vp = x.coeffRef( j - o, l );
        if( maxP < vp )
        {
          maxP = vp;
          maxL = l;

        } // fi

      } // rof
      lbl = inv_labels[ maxL ];
      p = maxP;
    }
    else
    {
      idx = Self::_ND( mIt->first, region );
      lbl = mIt->second;
      p = TScalar( 1 );
      mIt++;
      o++;

    } // fi
    out_labels->SetPixel( idx, lbl );
    out_probs->SetPixel( idx, p );

  } // rof
}

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalar >
unsigned long
fpa::Common::OriginalRandomWalker< _TImage, _TLabels, _TScalar >::
_1D( const TIndex& idx, const TRegion& region )
{
  unsigned long i = idx[ 0 ];
  unsigned long off = 1;
  typename TRegion::SizeType size = region.GetSize( );
  for( unsigned int d = 1; d < TIndex::Dimension; ++d )
  {
    off *= size[ d - 1 ];
    i += idx[ d ] * off;

  } // rof
  return( i );
}

// -------------------------------------------------------------------------
template< class _TImage, class _TLabels, class _TScalar >
typename fpa::Common::OriginalRandomWalker< _TImage, _TLabels, _TScalar >::
TIndex fpa::Common::OriginalRandomWalker< _TImage, _TLabels, _TScalar >::
_ND( const unsigned long& i, const TRegion& region )
{
  typename TRegion::SizeType size = region.GetSize( );

  unsigned long j = i;
  TIndex idx;
  if( TIndex::Dimension == 3 )
  {
    unsigned long z = size[ 0 ] * size[ 1 ];
    idx[ 2 ] = j / z;
    j -= idx[ 2 ] * z;

  } // fi
  idx[ 1 ] = j / size[ 0 ];
  idx[ 0 ] = j % size[ 0 ];
  return( idx );
}

#endif // __fpa__Common__OriginalRandomWalker__hxx__
// eof - $RCSfile$