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

#include "CTArteries.h"
#include "ui_CTArteries.h"
#include "ui_Parameters.h"
#include "FileDialog.h"

#include <set>

#include <itkImageFileReader.h>
#include <itkImageSeriesReader.h>
#include <itkQuadEdgeMesh.h>

#include <vtkMarchingCubes.h>
#include <vtkMarchingSquares.h>
#include <vtkPolyDataConnectivityFilter.h>
#include <vtkPolyDataMapper.h>
#include <vtkProperty.h>
#include <vtkRenderer.h>
#include <vtkRenderWindow.h>
#include <vtkSplineWidget.h>

#include <QtGui>

#include <ivq/ITK/CPRImageFilter.h>
#include <ivq/VTK/ImageActor.h>
#include <ivq/VTK/ImageViewer.h>
#include <ivq/VTK/MPRViewers.h>
#include <ivq/VTK/PolyDataActor.h>
#include <ivq/VTK/SeedWidgetOverImageActor.h>
#include <ivq/Qt/DicomSeriesSelectorDialog.h>
#include <ivq/Qt/RendererWidget.h>

#include "algorithms/RandomWalkSegmentation.h"

// -------------------------------------------------------------------------
#define QT_ACTION_CONN( x )                                             \
  this->connect(                                                        \
    this->m_UI->a##x, SIGNAL( triggered( ) ), this, SLOT( _s##x( ) )    \
    )

// -------------------------------------------------------------------------
class ContourCallback
  : public vtkCommand
{
public:
  static ContourCallback* New( )
    {
      return( new ContourCallback( ) );
    }
  virtual void Execute(
    vtkObject* caller, unsigned long eId, void* data
    ) override
    {
      ivq::VTK::ImageActor* actor =
        ivq::VTK::ImageActor::SafeDownCast( caller );
      if( eId == vtkCommand::InteractionEvent && actor != NULL )
      {
        ivq::VTK::ImageViewer* viewer = this->Window->m_UI->Slice->GetViewer( );
        if( this->Window->m_ContourActor != NULL )
          viewer->GetRenderer( )->RemoveViewProp( this->Window->m_ContourActor->GetActor( ) );

        double abnds[ 6 ];
        actor->GetActorBounds( abnds );
        unsigned int nSlice = actor->GetSliceNumber( );
        for( unsigned int i = 0; i < 16; ++i )
        {
          double w = double( 6.28 ) * double( i ) / double( 16 );
          CTArteries::TFourier::TPoint v = this->Window->m_Fourier[ nSlice ]( w );
          this->Window->m_Spline->SetHandlePosition(
            i, v[ 0 ], v[ 1 ], abnds[ 4 ]
            );

        } // rof
        this->Window->m_Spline->SetProjectionPosition(
          (
            actor->GetBounds( )[ 4 ] +
            actor->GetBounds( )[ 5 ]
            ) / double( 2 )
          );


        this->Window->m_ContourActor =
          this->Window->m_ContoursActors[ actor->GetSliceNumber( ) ];
        viewer->GetRenderer( )->AddViewProp( this->Window->m_ContourActor->GetActor( ) );
        viewer->Render( );

      } // fi
    }

protected:
  ContourCallback( )
    : vtkCommand( )
    {
    }
  virtual ~ContourCallback( )
    {
    }

private:
  // Purposely not implemented
  ContourCallback( const ContourCallback& other );
  ContourCallback& operator=( const ContourCallback& other );
public:
  CTArteries* Window;
};

// -------------------------------------------------------------------------
CTArteries::
CTArteries( int argc, char* argv[], QWidget* parent )
  : Superclass( parent ),
    m_UI( new Ui::CTArteries ),
    m_UIParameters( new Ui::Parameters )
{
  this->m_UI->setupUi( this );
  this->m_DlgParameters = new QDialog( );
  this->m_UIParameters->setupUi( this->m_DlgParameters );

  // Connect signals with slots
  QT_ACTION_CONN( Open );
  QT_ACTION_CONN( Config );
  QT_ACTION_CONN( Process );

  // Load log
  if( argc == 3 )
  {
    std::string cmd( argv[ 1 ] );
    std::string fname( argv[ 2 ] );
    if( cmd == "-log" )
      this->_ExecuteLog( fname );

  } // fi
}

// -------------------------------------------------------------------------
CTArteries::
~CTArteries( )
{
  delete this->m_UI;
  delete this->m_UIParameters;
  delete this->m_DlgParameters;
}

// -------------------------------------------------------------------------
template< class _TStrings >
void CTArteries::
_openImage( const _TStrings& fnames )
{
  // Save log
  std::ofstream str_log( "CTArteries.log" );
  str_log << fnames.size( ) << std::endl;
  for( const std::string& s: fnames )
    str_log << s << std::endl;
  str_log.close( );

  // Create appropriate reader
  itk::ImageSource< TImage >::Pointer reader;
  if( fnames.size( ) == 1 )
  {
    itk::ImageFileReader< TImage >::Pointer r =
      itk::ImageFileReader< TImage >::New( );
    r->SetFileName( *( fnames.begin( ) ) );
    reader = r;
  }
  else if( fnames.size( ) > 1 )
  {
    itk::ImageSeriesReader< TImage >::Pointer r =
      itk::ImageSeriesReader< TImage >::New( );
    typename _TStrings::const_iterator i;
    for( i = fnames.begin( ); i != fnames.end( ); ++i )
      r->AddFileName( *i );
    reader = r;

  } // fi

  // Execute reader
  if( reader.IsNull( ) )
    return;
  try
  {
    reader->Update( );
    this->m_Image = reader->GetOutput( );
    this->m_Image->DisconnectPipeline( );
  }
  catch( std::exception& err )
  {
    this->m_Image = NULL;
    QMessageBox::critical( this, "Error opening image", err.what( ) );

  } // yrt
  this->_showInputImage( );
}

// -------------------------------------------------------------------------
void CTArteries::
_openDicom( const std::string& dirname )
{
  ivq::Qt::DicomSeriesSelectorDialog dlg( this );
  dlg.setStartDir( dirname );
  if( dlg.exec( ) == QDialog::Accepted )
    this->_openImage( *( dlg.selectedFilenames( ) ) );
}

// -------------------------------------------------------------------------
void CTArteries::
_showInputImage( )
{
  if( this->m_Image.IsNotNull( ) )
  {
    this->m_VTKImage = TVTKImage::New( );
    this->m_VTKImage->SetInput( this->m_Image );
    this->m_VTKImage->Update( );

    ivq::VTK::MPRViewers* viewers = this->m_UI->MPR->GetViewers( );
    viewers->SetInputData( this->m_VTKImage->GetOutput( ) );
    viewers->SetColorWindow( 1000 );
    viewers->SetColorLevel( 100 );

    ivq::VTK::ImageViewer* view = viewers->GetView( 2 );
    this->m_Seeds =
      vtkSmartPointer< ivq::VTK::SeedWidgetOverImageActor >::New( );
    this->m_Seeds->SetActor( view->GetImageActor( ) );
    this->m_Seeds->SetInteractor( view->GetRenderWindow( )->GetInteractor( ) );

    viewers->Render( );
    viewers->ResetCameras( );
    this->m_Seeds->EnabledOn( );
    viewers->Render( );
  }
  else
  {
    this->m_VTKImage = NULL;

  } // fi
}

// -------------------------------------------------------------------------
void CTArteries::
_process( const std::vector< TImage::PointType >& seeds )
{
  // Save log
  std::ofstream str_log( "CTArteries.log", std::ios_base::app );
  str_log << this->m_UIParameters->Beta->value( ) << std::endl;
  str_log << this->m_UIParameters->Sigma->value( ) << std::endl;
  str_log << this->m_UIParameters->Radius->value( ) << std::endl;
  str_log << seeds.size( ) << std::endl;
  for( TImage::PointType seed: seeds )
    str_log << seed << std::endl;
  str_log.close( );

  // Create algorithm
  typedef RandomWalkSegmentation< TImage, TScalarImage > _TSegmentation;
  _TSegmentation::Pointer seg = _TSegmentation::New( );
  seg->SetInput( this->m_Image );
  seg->SetBeta( this->m_UIParameters->Beta->value( ) );
  seg->SetSigma( this->m_UIParameters->Sigma->value( ) );
  seg->SetRadius( this->m_UIParameters->Radius->value( ) );
  seg->SetStartSeed( seeds[ 0 ] );
  seg->SetEndSeed( seeds[ 1 ] );
  for( TImage::PointType seed: seeds )
    seg->AddSeed( seed );
  try
  {
    seg->DebugOn( );
    seg->Update( );
    seg->DebugOff( );
  }
  catch( std::exception& err )
  {
    QMessageBox::critical(
      this,
      QMessageBox::tr( "Error reading image" ),
      QMessageBox::tr( err.what( ) )
      );
    return;

  } // yrt
  this->m_Seeds->EnabledOff( );

  this->m_Segmentation = seg->GetOutput( );
  this->m_Axis = seg->GetOutputAxis( );
  this->m_Segmentation->DisconnectPipeline( );
  this->m_Axis->DisconnectPipeline( );
}

// -------------------------------------------------------------------------
void CTArteries::
_prepareQuantification( )
{
  this->m_Contours.clear( );
  this->m_ContoursActors.clear( );
  this->m_Fourier.clear( );

  TImage::RegionType ext_region = this->m_CPRImage->GetRequestedRegion( );
  for( long z = 0; z < ext_region.GetSize( )[ 2 ]; ++z )
  {
    TImage::IndexType slc_idx = ext_region.GetIndex( );
    TImage::SizeType slc_size = ext_region.GetSize( );
    slc_idx[ 2 ] += z;
    slc_size[ 2 ] = 1;
    TImage::RegionType slc_region;
    slc_region.SetIndex( slc_idx );
    slc_region.SetSize( slc_size );

    typedef itk::RegionOfInterestImageFilter< TScalarImage, TScalarImage > _TROI;
    _TROI::Pointer roi = _TROI::New( );
    roi->SetInput( this->m_CPRSegmentation );
    roi->SetRegionOfInterest( slc_region );

    /* TODO
       slc_idx[ 2 ] = 1;
       slc_size[ 2 ] = 0;
       slc_region.SetIndex( slc_idx );
       slc_region.SetSize( slc_size );
       typedef itk::Image< TBinPixel, 2 > _TBinSlice;
       typedef itk::ExtractImageFilter< TBinImage, _TBinSlice > _TCollapse;
       _TCollapse::Pointer collapse = _TCollapse::New( );
       collapse->SetInput( roi->GetOutput( ) );
       collapse->SetDirectionCollapseToIdentity( );
       collapse->SetExtractionRegion( slc_region );
       collapse->UpdateLargestPossibleRegion( );
    */
    typedef TScalarImage _TScalarSlice;

    typedef itk::ImageToVTKImageFilter< _TScalarSlice > _TVTKScalarSlice;
    _TVTKScalarSlice::Pointer vtk_slice = _TVTKScalarSlice::New( );
    vtk_slice->SetInput( roi->GetOutput( ) );
    vtk_slice->Update( );

    double r[ 2 ];
    vtk_slice->GetOutput( )->GetScalarRange( r );

    vtkSmartPointer< vtkMarchingSquares > ms =
      vtkSmartPointer< vtkMarchingSquares >::New( );
    ms->SetInputData( vtk_slice->GetOutput( ) );
    ms->SetValue( 0, 0 );

    vtkSmartPointer< vtkPolyDataConnectivityFilter > conn =
      vtkSmartPointer< vtkPolyDataConnectivityFilter >::New( );
    conn->SetInputConnection( ms->GetOutputPort( ) );
    conn->SetExtractionModeToClosestPointRegion( );
    conn->SetClosestPoint( 0, 0, 0 );
    conn->Update( );

    vtkSmartPointer< vtkPolyData > cnt = vtkSmartPointer< vtkPolyData >::New( );
    cnt->DeepCopy( conn->GetOutput( ) );
    this->m_Contours.push_back( cnt );

    typedef itk::QuadEdgeMesh< double, 2 > _TQEMesh;
    typedef _TQEMesh::PointType _T2DPoint;
    typedef std::vector< _T2DPoint > _T2DPoints;
    _T2DPoints points;
    _TQEMesh::Pointer mesh = _TQEMesh::New( );
    for( unsigned int i = 0; i < cnt->GetNumberOfPoints( ); ++i )
    {
      _T2DPoint pnt;
      double p[ 3 ];
      cnt->GetPoint( i, p );
      pnt[ 0 ] = p[ 0 ];
      pnt[ 1 ] = p[ 1 ];
      mesh->AddPoint( pnt );

    } // rof
    vtkCellArray* lines = cnt->GetLines( );
    lines->InitTraversal( );
    vtkIdType* ids = new vtkIdType[ 2 ];
    vtkIdType npts;
    while( lines->GetNextCell( npts, ids ) != 0 )
      mesh->AddEdge( ids[ 0 ], ids[ 1 ] );
    delete ids;
    _TQEMesh::QEPrimal* edge = mesh->GetEdge( );
    if( edge != NULL )
    {
      mesh->AddFace( edge );
      edge = mesh->GetEdge( );
      for( auto eIt = edge->BeginGeomLnext( ); eIt != edge->EndGeomLnext( ); ++eIt )
        points.push_back( mesh->GetPoint( *eIt ) );

    } // fi

    TFourier f( points.begin( ), points.end( ), 6 );
    f.SetOrderingToCounterClockWise( );
    this->m_Fourier.push_back( f );

    vtkSmartPointer< ivq::VTK::PolyDataActor > actor =
      vtkSmartPointer< ivq::VTK::PolyDataActor >::New( );
    actor->SetInputData( cnt );
    actor->GetMapper( )->ScalarVisibilityOff( );
    actor->Update( );
    actor->GetProperty( )->SetColor( 1, 0, 0 );
    actor->GetProperty( )->SetLineWidth( 2 );
    this->m_ContoursActors.push_back( actor );

  } // rof

  ivq::VTK::ImageViewer* slice_view = this->m_UI->Slice->GetViewer( );
  vtkSmartPointer< ContourCallback > cbk =
    vtkSmartPointer< ContourCallback >::New( );
  cbk->Window = this;
  slice_view->GetImageActor( )->AddObserver( vtkCommand::InteractionEvent, cbk );
  this->m_ContourActor = NULL;

  this->m_Spline = vtkSmartPointer< vtkSplineWidget >::New( );
  this->m_Spline->SetCurrentRenderer( slice_view->GetRenderer( ) );
  this->m_Spline->SetDefaultRenderer( slice_view->GetRenderer( ) );
  this->m_Spline->SetInputData( slice_view->GetInput( ) );
  this->m_Spline->SetProp3D( slice_view->GetImageActor( ) );
  this->m_Spline->SetInteractor( slice_view->GetRenderWindow( )->GetInteractor( ) );
  double bnds[ 6 ];
  slice_view->GetInput( )->GetBounds( bnds );
  this->m_Spline->PlaceWidget(
    bnds[ 0 ], bnds[ 1 ],
    bnds[ 2 ], bnds[ 3 ],
    bnds[ 4 ], bnds[ 5 ]
    );
  this->m_Spline->ProjectToPlaneOn( );
  this->m_Spline->SetProjectionNormalToZAxes( );
  this->m_Spline->SetProjectionPosition(
    (
      slice_view->GetImageActor( )->GetBounds( )[ 4 ] +
      slice_view->GetImageActor( )->GetBounds( )[ 5 ]
      ) / double( 2 )
    );
  this->m_Spline->SetHandleSize( 0.005 );
  this->m_Spline->ClosedOn( );
  this->m_Spline->SetNumberOfHandles( 16 );
  this->m_Spline->EnabledOn( );

  QVector< double > xdata, dmindata, dmaxdata, areadata;
  double data[ 3 ];
  for( unsigned int i = 0; i < this->m_Fourier.size( ); ++i )
  {
    double a, b, t, p;
    this->m_Fourier[ i ].GetEllipse( 1, a, b, t, p );
    double area = this->m_Fourier[ i ].GetArea( );

    xdata.push_back( double( i ) );
    dmindata.push_back( b );
    dmaxdata.push_back( a );
    areadata.push_back( area );

  } // rof
  this->m_UI->Plot->addGraph( this->m_UI->Plot->xAxis, this->m_UI->Plot->yAxis );
  this->m_UI->Plot->graph( 0 )->setPen( QPen( QColor( 255, 0, 0 ) ) );
  this->m_UI->Plot->graph( 0 )->setName( "Minimum diameter (mm)" );
  this->m_UI->Plot->graph( 0 )->setData( xdata, dmindata );

  this->m_UI->Plot->addGraph( this->m_UI->Plot->xAxis, this->m_UI->Plot->yAxis );
  this->m_UI->Plot->graph( 1 )->setPen( QPen( QColor( 0, 255, 0 ) ) );
  this->m_UI->Plot->graph( 1 )->setName( "Maximum diameter (mm)" );
  this->m_UI->Plot->graph( 1 )->setData( xdata, dmaxdata );

  this->m_UI->Plot->addGraph( this->m_UI->Plot->xAxis, this->m_UI->Plot->yAxis2 );
  this->m_UI->Plot->graph( 2 )->setPen( QPen( QColor( 0, 0, 255 ) ) );
  this->m_UI->Plot->graph( 2 )->setName( "Area (mm2)" );
  this->m_UI->Plot->graph( 2 )->setData( xdata, areadata );

  this->m_UI->Plot->xAxis->setLabel( "Position" );
  this->m_UI->Plot->yAxis->setLabel( "Diameters (mm)" );
  this->m_UI->Plot->yAxis2->setLabel( "Areas (mm2)" );

  this->m_UI->Plot->legend->setVisible( true );
  this->m_UI->Plot->xAxis->setVisible( true );
  this->m_UI->Plot->yAxis->setVisible( true );
  this->m_UI->Plot->yAxis2->setVisible( true );

  this->m_UI->Plot->rescaleAxes( );
  this->m_UI->Plot->replot( );
}

// -------------------------------------------------------------------------
void CTArteries::
_showProcessResults( )
{
  if( this->m_Segmentation.IsNull( ) || this->m_Axis.IsNull( ) )
    return;

  this->m_VTKSegmentation = TVTKScalarImage::New( );
  this->m_VTKSegmentation->SetInput( this->m_Segmentation );
  this->m_VTKSegmentation->Update( );

  // Show surface
  vtkImageData* seg = this->m_VTKSegmentation->GetOutput( );
  double r[ 2 ];
  seg->GetScalarRange( r );
  vtkSmartPointer< vtkMarchingCubes > mc =
    vtkSmartPointer< vtkMarchingCubes >::New( );
  mc->SetInputData( seg );
  mc->SetValue( 0, 0 );
  mc->Update( );
  this->m_Surface = vtkSmartPointer< ivq::VTK::PolyDataActor >::New( );
  this->m_Surface->SetInputConnection( mc->GetOutputPort( ) );
  this->m_Surface->Update( );

  // Prepare curve
  this->m_Curve = TCurve::New( );
  for( unsigned int i = 0; i < this->m_Axis->GetSize( ); ++i )
    this->m_Curve->AddPoint( this->m_Axis->GetPoint( i ) );
  this->m_Curve->Smooth( 2 );

  // Compute CPRs
  this->_CPR(
    this->m_CPRImage, this->m_VTKCPRImage,
    this->m_Image, this->m_Curve
    );
  this->_CPR(
    this->m_CPRSegmentation, this->m_VTKCPRSegmentation,
    this->m_Segmentation, this->m_Curve
    );

  // Show all data
  vtkRenderer* renderer =
    this->m_UI->MPR->GetRendererWidget( )->GetRenderer( );
  renderer->AddViewProp( this->m_Surface->GetActor( ) );
  renderer->Render( );
  renderer->ResetCamera( );
  this->m_UI->MPR->GetRendererWidget( )->GetRenderWindow( )->Render( );

  ivq::VTK::ImageViewer* cpr_view = this->m_UI->CPR->GetViewer( );
  ivq::VTK::ImageViewer* slice_view = this->m_UI->Slice->GetViewer( );

  cpr_view->SetInputData( this->m_VTKCPRImage->GetOutput( ) );
  cpr_view->SetColorWindow( 1000 );
  cpr_view->SetColorLevel( 100 );
  cpr_view->SetSliceOrientationToXZ( );
  cpr_view->Render( );

  slice_view->SetInputData( this->m_VTKCPRImage->GetOutput( ) );
  slice_view->SetColorWindow( 1000 );
  slice_view->SetColorLevel( 100 );
  slice_view->SetSliceOrientationToXY( );
  slice_view->Render( );

  this->_prepareQuantification( );
}

// -------------------------------------------------------------------------
void CTArteries::
_ExecuteLog( const std::string& fname )
{
  std::ifstream str_log( fname.c_str( ) );
  if( str_log )
  {
    str_log.close( );

  } // fi
}

// -------------------------------------------------------------------------
void CTArteries::
_sOpen( )
{
  FileDialog dlg( this );
  dlg.setNameFilter(
    "Medical images (*.mhd *.dcm);;Other images (*.png *.jpg *.bmp);;\
All files (*)"
    );
  if( dlg.exec( ) == QDialog::Accepted )
  {
    std::set< std::string > filenames, dirnames;
    QStringList names = dlg.selectedFiles( );
    for( QString name: names )
    {
      QFileInfo info( name );
      if( info.isDir( ) )
        dirnames.insert( name.toStdString( ) );
      else
        filenames.insert( name.toStdString( ) );

    } // rof

    if( dirnames.size( ) == 1 && filenames.size( ) == 0 )
      this->_openDicom( *( dirnames.begin( ) ) );
    else if( dirnames.size( ) == 0 && filenames.size( ) > 0 )
      this->_openImage( filenames );
    else
      QMessageBox::critical(
        this,
        "Error opening image",
        "Directory - files mixed up, don't know what to do."
        );

  } // fi
}

// -------------------------------------------------------------------------
void CTArteries::
_sConfig( )
{
  this->m_DlgParameters->exec( );
}

// -------------------------------------------------------------------------
void CTArteries::
_sProcess( )
{
  // Get seeds
  typedef ivq::VTK::SeedWidgetOverImageActor::TSeeds _TSeeds;
  if( this->m_Image.IsNull( ) || this->m_Seeds.GetPointer( ) == NULL )
  {
    QMessageBox::critical( this, "Error processing", "No valid input image." );
    return;

  } // fi
  std::vector< TImage::PointType > seeds;
  for( _TSeeds::value_type sValue: this->m_Seeds->GetSeeds( ) )
  {
    for( unsigned int i = 0; i < sValue.second.size( ); i += 3 )
    {
      TImage::PointType pnt;
      pnt[ 0 ] = sValue.second[ i + 0 ];
      pnt[ 1 ] = sValue.second[ i + 1 ];
      pnt[ 2 ] = sValue.second[ i + 2 ];
      seeds.push_back( pnt );

    } // rof

  } // rof
  if( seeds.size( ) < 2 )
  {
    QMessageBox::critical( this, "Error processing", "Not enough seeds." );
    return;

  } // fi
  this->_process( seeds );
  this->_showProcessResults( );
}

// -------------------------------------------------------------------------
template< class _TImagePtr, class _TVTKImagePtr, class _TCurvePtr >
void CTArteries::
_CPR(
  _TImagePtr& output, _TVTKImagePtr& vtk_output,
  const _TImagePtr& input, const _TCurvePtr& curve
  )
{
  typedef typename _TImagePtr::ObjectType _TImage;
  typedef typename _TCurvePtr::ObjectType _TCurve;
  typedef ivq::ITK::CPRImageFilter< _TImage, _TCurve > _TFilter;

  typename _TFilter::Pointer cpr = _TFilter::New( );
  cpr->SetInput( input );
  cpr->SetCurve( curve );
  cpr->SetSliceRadius( 40 );
  cpr->Update( );
  output = cpr->GetOutput( );
  output->DisconnectPipeline( );

  vtk_output = _TVTKImagePtr::ObjectType::New( );
  vtk_output->SetInput( output );
  vtk_output->Update( );
}

// -------------------------------------------------------------------------
#include <QApplication>

int main( int argc, char* argv[] )
{
  QApplication a( argc, argv );
  CTArteries w( argc, argv );
  w.show( );
  return( a.exec( ) );
}

// eof - $RCSfile$