
#ifndef _LinearTetrahedralMesh_h
#define _LinearTetrahedralMesh_h

#include "itkImage.h"
#include "itkVector.h"
#include "itkListSample.h"
#include "itkKdTree.h"

#include "vtkSmartPointer.h"
#include "vtkUnstructuredGrid.h"

#include "vnl/vnl_matrix.h"
#include "vnl/vnl_vector.h"
#include "vnl/algo/vnl_matrix_inverse.h"
#include "vnl/algo/vnl_qr.h"

#include "DynArray.h"

class LinearTetrahedralMesh
{

public:

  typedef float ScalarType;

  typedef vnl_matrix<ScalarType> MatrixType;
  typedef vnl_vector<ScalarType> VectorType;

  typedef vnl_matrix_inverse<ScalarType> MatrixInverseType;
  typedef vnl_qr<ScalarType> MatrixQRType;

  typedef itk::Vector<float, 3> PointType;
  typedef itk::Statistics::ListSample<PointType> SampleType;
  typedef itk::Statistics::KdTree<SampleType> KdTreeType;
  typedef KdTreeType::InstanceIdentifierVectorType PointIDVectorType;

  LinearTetrahedralMesh();
  virtual ~LinearTetrahedralMesh();

  void ReadVTKFile(const char* fn);

  inline void SetVTKMesh(vtkUnstructuredGrid* mesh)
  {
    vtkSmartPointer<vtkUnstructuredGrid> ptr = mesh;
    this->SetVTKMesh(ptr);
  }
  inline void SetVTKMesh(const vtkSmartPointer<vtkUnstructuredGrid>& mesh)
  {
    m_VTKMesh = mesh;
    this->RecomputeMappings();
    // Reset Kd tree, build it only if neighbor query is made
    m_KdTree = 0;
    m_KdTreeSample = 0;
  }

  vtkUnstructuredGrid* GetVTKMesh()
  { return m_VTKMesh.GetPointer(); }

  VectorType ComputeShapeFunctions(unsigned int el, const VectorType& x);

  MatrixType ComputeShapeFunctionDerivatives(unsigned int el);

  float ComputeElementVolume(unsigned int el);

  inline unsigned int GetNumberOfElements()
  { 
    if (m_VTKMesh.GetPointer() == 0)
      return 0;
    return m_VTKMesh->GetNumberOfCells();
  }

  inline unsigned int GetNumberOfPoints()
  { 
    if (m_VTKMesh.GetPointer() == 0)
      return 0;
    return m_VTKMesh->GetNumberOfPoints();
  }

  PointIDVectorType FindClosestPoints(PointType& query, unsigned int n);

  void ClearMappings();

protected:
  void RecomputeMappings();

private:

  vtkSmartPointer<vtkUnstructuredGrid> m_VTKMesh;

  // List of matrices that map global coords to shape functions
  DynArray<MatrixType> m_ElementWeightMappings;

  DynArray<float> m_ElementVolumes;

  KdTreeType::Pointer m_KdTree;

  // Have to explicitly keep track of kd tree sample (ITK bug?)
  SampleType::Pointer m_KdTreeSample;

};

#endif
