
#include "itkImage.h"
#include "itkImageFileReader.h"
#include "itkImageRegionIterator.h"

#include "itkOutputWindow.h"
#include "itkTextOutput.h"

#include "vtkOutputWindow.h"
#include "vtkSmartPointer.h"
#include "vtkUnstructuredGrid.h"
#include "vtkUnstructuredGridWriter.h"

#include "vtkTextOutputWindow.h"

#include "DynArray.h"
#include "Timer.h"

#include "BrainMeshGenerator.h"

int main(int argc, char** argv)
{
  if (argc != 5)
  {
    std::cerr << "Usage: " << argv[0] << " <label> <spacing> <iters> <dt>" << std::endl;
    std::cerr << "Label image contains wm, gm, ventricles, and falx" << std::endl;
    return -1;
  }

  Timer timer;
  timer.Start();

  double spacing = atof(argv[2]);
  if (spacing < 0.01)
  {
    std::cerr << "Spacing should be >> 0" << std::endl;
    return -1;
  }

  unsigned int iters = (unsigned int)atoi(argv[3]);

  double dt = atof(argv[4]);
  if (dt < 1e-10 || dt > 0.625)
  {
    std::cerr << "Time step should be > 0 and < 0.625" << std::endl;
    return -1;
  }

  itk::OutputWindow::SetInstance(itk::TextOutput::New());
  vtkOutputWindow::SetInstance(vtkTextOutputWindow::New());

  typedef itk::Image<unsigned char, 3> ByteImageType;
  
  // Read label image
  typedef itk::ImageFileReader<ByteImageType> ReaderType;
  
  ByteImageType::Pointer labelImg;
  {
    ReaderType::Pointer reader = ReaderType::New();
    reader->SetFileName(argv[1]);
    reader->Update();

    labelImg = reader->GetOutput();
  }

  // Exclude ventricles from mesh
  typedef itk::ImageRegionIterator<ByteImageType> IteratorType;
  IteratorType labelIt(labelImg, labelImg->GetLargestPossibleRegion());

  labelIt.GoToBegin();
  while (!labelIt.IsAtEnd())
  {
    unsigned char c = labelIt.Get();
    if (c == 3)
      labelIt.Set(0);
    ++labelIt;
  }

  // Generate tetrahedral mesh
  BrainMeshGenerator meshgen;

  meshgen.SetLabelImage(labelImg);
  meshgen.SetInitialSpacing(spacing);
  meshgen.SetMaximumIterations(iters);
  meshgen.SetTimeStep(dt);

  vtkSmartPointer<vtkUnstructuredGrid> mesh = meshgen.GenerateMesh();

  // Write mesh
  vtkSmartPointer<vtkUnstructuredGridWriter> meshWriter = 
    vtkSmartPointer<vtkUnstructuredGridWriter>::New();

  meshWriter->SetFileTypeToBinary();
  meshWriter->SetFileName("mesh_gen_out.vtk");
  meshWriter->SetInput(mesh);
  meshWriter->Update();

  timer.Stop();
  std::cout << "Mesh generation took "
    << timer.GetElapsedHours() << " hours, "
    << timer.GetElapsedMinutes() << " minutes, "
    << timer.GetElapsedSeconds() << " seconds"
    << std::endl;

  return 0;

}
