//-------Surface extraction from a segmented object and mesh(the extracted surface) visualization 

#include "itkImageFileReader.h"
#include "itkBinaryMask3DMeshSource.h"
#include "itkImage.h"
#include "itkMesh.h"

#include "itkQuadEdgeMesh.h"
#include "itkVTKPolyDataReader.h"
#include "itkQuadEdgeMeshExtendedTraits.h"
#include "itkQuadEdgeMeshScalarDataVTKPolyDataWriter.h"

#include "itkQuadEdgeMeshDiscreteGaussianCurvatureEstimator.h"
#include "itkQuadEdgeMeshDiscreteMeanCurvatureEstimator.h"
#include "itkQuadEdgeMeshDiscreteMinCurvatureEstimator.h"
#include "itkQuadEdgeMeshDiscreteMaxCurvatureEstimator.h"
#include "itkQuadEdgeMeshScalarDataVTKPolyDataWriter.h"

#include <iostream>
#include "itkTriangleCell.h"
#include "itkQuadrilateralCell.h"
#include "vtkDataSetWriter.h"
#include "vtkDataSetMapper.h"
#include "vtkRenderer.h"
#include "vtkRenderWindow.h"
#include "vtkActor.h"
#include "vtkRenderWindowInteractor.h"
#include "vtkDataSetReader.h"
#include "vtkUnstructuredGrid.h"
#include "vtkDataSet.h"
#include "vtkCellArray.h"

using namespace itk;

typedef double CoordType;
typedef QuadEdgeMeshExtendedTraits < CoordType, 3,  2, CoordType, CoordType,  CoordType,  bool, bool > Traits;
typedef QuadEdgeMesh< CoordType, 3, Traits > MeshType;
//typedef itk::Mesh<double, 3,itk::DefaultStaticMeshTraits< double, 3, 3, double, double > > MeshType;

void Display(vtkUnstructuredGrid* itkgrid)
{
  vtkRenderer* ren1 = vtkRenderer::New();
  vtkRenderWindow* renWin = vtkRenderWindow::New();
  renWin->AddRenderer(ren1);
  vtkRenderWindowInteractor* inter = vtkRenderWindowInteractor::New();
  inter->SetRenderWindow(renWin);

  vtkDataSetMapper* mapper = vtkDataSetMapper::New();
  mapper->SetInput(itkgrid);
  vtkActor* actor = vtkActor::New();
  actor->SetMapper(mapper);
  ren1->AddActor(actor);
  renWin->Render();
  inter->Start();
  
  mapper->Delete();
  actor->Delete();
  ren1->Delete();
  renWin->Delete();
}


class VistVTKCellsClass
{
  vtkCellArray* m_Cells;
  int* m_LastCell;
  int* m_TypeArray;
public:
  // typedef the itk cells we are interested in
  typedef itk::CellInterface<MeshType::PixelType,MeshType::CellTraits >  CellInterfaceType;
  typedef itk::TriangleCell<CellInterfaceType>      floatTriangleCell;
  typedef itk::QuadrilateralCell<CellInterfaceType> floatQuadrilateralCell;

  // Set the vtkCellArray that will be constructed
  void SetCellArray(vtkCellArray* a) 
    {
      m_Cells = a;
    }
  // Set the cell counter pointer
  void SetCellCounter(int* i)
    {
      m_LastCell = i;
    }
  // Set the type array for storing the vtk cell types
  void SetTypeArray(int* i)
    {
      m_TypeArray = i;
    }
  // Visit a triangle and create the VTK_TRIANGLE cell 
  void Visit(unsigned long, floatTriangleCell* t)
    {
      m_Cells->InsertNextCell(3,  (vtkIdType*)t->PointIdsBegin());
      m_TypeArray[*m_LastCell] = VTK_TRIANGLE;
      (*m_LastCell)++;
    }
  // Visit a triangle and create the VTK_QUAD cell 
  void Visit(unsigned long, floatQuadrilateralCell* t)
    {
      m_Cells->InsertNextCell(4,  (vtkIdType*)t->PointIdsBegin());
      m_TypeArray[*m_LastCell] = VTK_QUAD;
      (*m_LastCell)++;
    }
};
  
typedef itk::CellInterfaceVisitorImplementation<vtkFloatingPointType, MeshType::CellTraits,itk::TriangleCell< itk::CellInterface<MeshType::PixelType, MeshType::CellTraits > >, 
  VistVTKCellsClass> TriangleVisitor;

typedef itk::CellInterfaceVisitorImplementation<vtkFloatingPointType, MeshType::CellTraits,itk::QuadrilateralCell< itk::CellInterface<MeshType::PixelType, MeshType::CellTraits > >, 
  VistVTKCellsClass> QuadrilateralVisitor;

vtkUnstructuredGrid* MeshToUnstructuredGrid(MeshType* mesh)
{
  // Get the number of points in the mesh
  int numPoints = mesh->GetNumberOfPoints();
  if(numPoints == 0)
    {
    mesh->Print(std::cerr);
    std::cerr << "no points in Grid " << std::endl;
    exit(-1);
    }
  // Create a vtkUnstructuredGrid
  vtkUnstructuredGrid* vgrid = vtkUnstructuredGrid::New();

  // Create the vtkPoints object and set the number of points
  vtkPoints* vpoints = vtkPoints::New();
  vpoints->SetNumberOfPoints(numPoints);
  // iterate over all the points in the itk mesh filling in
  // the vtkPoints object as we go
  MeshType::PointsContainer::Pointer points = mesh->GetPoints();
  for(MeshType::PointsContainer::Iterator i = points->Begin();
      i != points->End(); ++i)
    {
    // Get the point index from the point container iterator
    int idx = i->Index();
    // Set the vtk point at the index with the the coord array from itk
    // itk returns a const pointer, but vtk is not const correct, so
    // we have to use a const cast to get rid of the const
    vpoints->SetPoint(idx, const_cast<vtkFloatingPointType*>(i->Value().GetDataPointer()));
    }
  // Set the points on the vtk grid
  vgrid->SetPoints(vpoints);
  
  // Now create the cells using the MulitVisitor
  // 1. Create a MultiVisitor
  MeshType::CellType::MultiVisitor::Pointer mv = MeshType::CellType::MultiVisitor::New();
  // 2. Create a triangle and quadrilateral visitor
  TriangleVisitor::Pointer tv = TriangleVisitor::New();
  QuadrilateralVisitor::Pointer qv =  QuadrilateralVisitor::New();
  // 3. Set up the visitors
  int vtkCellCount = 0; // running counter for current cell being inserted into vtk
  int numCells = mesh->GetNumberOfCells();
  int *types = new int[numCells]; // type array for vtk 
  // create vtk cells and estimate the size
  vtkCellArray* cells = vtkCellArray::New();
  cells->EstimateSize(numCells, 4);
  // Set the TypeArray CellCount and CellArray for both visitors
  tv->SetTypeArray(types);
  tv->SetCellCounter(&vtkCellCount);
  tv->SetCellArray(cells);
  qv->SetTypeArray(types);
  qv->SetCellCounter(&vtkCellCount);
  qv->SetCellArray(cells);
  // add the visitors to the multivisitor
  mv->AddVisitor(tv);
  mv->AddVisitor(qv);
  // Now ask the mesh to accept the multivisitor which
  // will Call Visit for each cell in the mesh that matches the
  // cell types of the visitors added to the MultiVisitor
  mesh->Accept(mv);
  
  // Now set the cells on the vtk grid with the type array and cell array
  vgrid->SetCells(types, cells);
  
  // Clean up vtk objects (no vtkSmartPointer ... )
  cells->Delete();
  vpoints->Delete();
  // return the vtkUnstructuredGrid
  return vgrid;
}

int main(void ) 
{ 
//1---------- surface extraction from a segmented 3D object
  const unsigned int Dimension = 3;
  typedef float  PixelType; 
  typedef itk::Image< PixelType, Dimension >   ImageType;
  typedef itk::ImageFileReader< ImageType >    ReaderType;
  ReaderType::Pointer reader = ReaderType::New();

  typedef itk::BinaryMask3DMeshSource< ImageType, MeshType >   MeshSourceType;
  MeshSourceType::Pointer meshSource = MeshSourceType::New();

  reader->SetFileName("C:/Images/ConfiConnect_brainweb165a10f17_whiteMatter.hdr");	// a segmented volume

  try
    {
    reader->Update();
    }
  catch( itk::ExceptionObject & exp )
    {
    std::cerr << "Exception thrown while reading the input file " << std::endl;
    std::cerr << exp << std::endl;
    return EXIT_FAILURE;
    }

  const PixelType objectValue = static_cast<PixelType>( 255); 
  meshSource->SetObjectValue( objectValue );
  meshSource->SetInput( reader->GetOutput() );

 
  try
    {
    meshSource->Update();
    }
  catch( itk::ExceptionObject & exp )
    {
    std::cerr << "Exception thrown during Update() " << std::endl;
    std::cerr << exp << std::endl;
    return EXIT_FAILURE;
    }

  std::cout << "Nodes = " << meshSource->GetNumberOfNodes() << std::endl;
  std::cout << "Cells = " << meshSource->GetNumberOfCells() << std::endl;
  
   //2---------- Mesh (extracted surface)visualization
  vtkUnstructuredGrid* grid = MeshToUnstructuredGrid(meshSource->GetOutput()); // display the initial sphere
  Display(grid);
   
  //3----------------Curvature estimation

  return EXIT_SUCCESS;
}