
////////////////////////////////////////////////////////////////////////////////
//
// Simulate tumor mass effect using a linear elastic biomechanical model
//
////////////////////////////////////////////////////////////////////////////////

// prastawa@cs.unc.edu 4/2006

#ifndef _TumorMassEffectGenerator_h
#define _TumorMassEffectGenerator_h

#include "itkImage.h"

#include "itkMultiThreader.h"
#include "itkSimpleFastMutexLock.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 "vnl/algo/vnl_svd.h"

#include "LinearTetrahedralMesh.h"

#include <vector>

class TumorMassEffectGenerator
{

public:

  typedef itk::Image<unsigned char, 3> ByteImageType;
  typedef ByteImageType::Pointer ByteImagePointer;
  typedef ByteImageType::IndexType ByteImageIndexType;
  typedef ByteImageType::PointType ByteImagePointType;
  typedef ByteImageType::OffsetType ByteImageOffsetType;
  typedef ByteImageType::RegionType ByteImageRegionType;
  typedef ByteImageType::SizeType ByteImageSizeType;
  typedef ByteImageType::SpacingType ByteImageSpacingType;

  typedef itk::Image<float, 3> FloatImageType;
  typedef FloatImageType::Pointer FloatImagePointer;
  typedef FloatImageType::IndexType FloatImageIndexType;
  typedef FloatImageType::PointType FloatImagePointType;
  typedef FloatImageType::OffsetType FloatImageOffsetType;
  typedef FloatImageType::RegionType FloatImageRegionType;
  typedef FloatImageType::SizeType FloatImageSizeType;
  typedef FloatImageType::SpacingType FloatImageSpacingType;

  typedef itk::Vector<float, 3> DisplacementType;
  typedef itk::Image<DisplacementType, 3> DeformationFieldType;
  typedef DeformationFieldType::Pointer DeformationFieldPointer;

  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 vnl_svd<ScalarType> MatrixSVDType;

  typedef enum{Internal, SlidingBC, FixedBC, PressureBC} BCEnumType;

  typedef struct
  {
    unsigned int ids[3];
  }
  TriangleSurface;

  TumorMassEffectGenerator();
  virtual ~TumorMassEffectGenerator();

  void SetLabelImage(ByteImagePointer labelImg);

  inline void SetDeformationIterations(unsigned int n)
  { m_DeformationIterations = n; }

  inline void SetPressure(float p)
  { m_Pressure = p; }

  inline void SetInitialVTKMesh(vtkUnstructuredGrid* mesh)
  {
    m_InitialVTKMesh = mesh;
  }

  inline void SetVMFKappa(float d)
  { m_VMFKappa = d; }

  inline float GetVMFKappa() const { return m_VMFKappa; }

  LinearTetrahedralMesh* GetCurrentMesh() const
  { return m_CurrentMesh; }

  // Set the Young's modulus and Poisson ratio for the brain tissue and falx
  void SetBrainMaterialParameters(float E, float nu);
  void SetFalxMaterialParameters(float E, float nu);

  void ComputeDeformation();

  DeformationFieldPointer GetDeformation() { return m_Deformation; }
  DeformationFieldPointer GetInverseDeformation() { return m_InverseDeformation; }

  void SetUseQHull(bool b) { m_UseQHull = b; }
  bool GetUseQHull() const { return m_UseQHull; }

  void SetDamping(float f) { m_Damping = f; }
  float GetDamping() const { return m_Damping; }

  inline void SetDeformationSolverIterations(unsigned int n)
  { m_DeformationSolverIterations = n; }

  inline void SetNumberOfThreads(unsigned int n)
  { m_NumberOfThreads = n; }
  inline unsigned int GetNumberOfThreads() const { return m_NumberOfThreads; }

  unsigned int GetNextElement();

  inline void ResetElementCounter() { m_ElementCounter = 0; }

protected:

  // Remove points inside tumor and add tumor boundary points to the mesh
  virtual void ModifyInitialMesh();

  // Compute solution for a given location as linear combination of element
  // shape functions and node solutions
  VectorType ComputeSolution(const VectorType& x);

  void ComputeTumorCentroids();

  virtual void DetermineBCs();

  MatrixType BuildElasticityMatrix(float E, float nu);

  // Initialize node solutions with the forces (can be either surface traction
  // or body forces)
  virtual void ComputeForces(MatrixType& solutions);

  void ComputeElementLabels();
  void DetermineSingularPoints();

  VectorType Precondition(const VectorType& x);

  VectorType GMRES(const VectorType& init, const VectorType& f, unsigned int maxiters);

  // Compute the solution
  void UpdateSolutionsJacobi();
  void RefineSolutionsGMRES();

  //VectorType ProductKu(const MatrixType& disp);
  VectorType ProductKu(const VectorType& u);
  VectorType ProductKuNoDiagonal(const VectorType& u);

  // Displace points in mesh
  void UpdateMesh(DeformationFieldPointer mapImg);

  // Update deformation field using solution for current mesh
  void UpdateDeformationField(DeformationFieldPointer defImg);

  // Invert current deformation and accumulate
  void InvertAccumulateDeformationField(DeformationFieldPointer invDefImg, DeformationFieldPointer defImg, DeformationFieldPointer currDefImg);

  // (x + v1) \circ (x + v2)
  DeformationFieldPointer Compose(DeformationFieldPointer inDef1, DeformationFieldPointer inDef2);

protected:

  float m_PointTolerance;

  bool m_UseQHull;

  vtkSmartPointer<vtkUnstructuredGrid> m_InitialVTKMesh;

  LinearTetrahedralMesh* m_InitialMesh;
  LinearTetrahedralMesh* m_CurrentMesh;

  MatrixType m_ElasticityMatrix;

  DynArray<unsigned char> m_NodeLabels;

  DynArray<unsigned char> m_ElementLabels;

  DynArray<unsigned char> m_SingularPoints;

  std::vector<BCEnumType> m_NodeBCMarkers;

  ByteImagePointer m_LabelImage;

  // For computing tumor expansion forces
  float m_Pressure;
  //VectorType m_TumorCentroid;
  DynArray<VectorType> m_TumorCentroids;
  DynArray<TriangleSurface> m_TumorPressureSurfaces;

  FloatImagePointer m_BrainDistanceImage;

  unsigned int m_DeformationIterations;

  MatrixType m_NodeSolutions; // Solutions for an iteration
  MatrixType m_NodeCompSolutions; // Solutions compounded over current iters

  MatrixType m_BrainElasticityMatrix;
  MatrixType m_FalxElasticityMatrix;
  MatrixType m_DuraElasticityMatrix;

  float m_VMFKappa;

  float m_Damping;

  unsigned int m_DeformationSolverIterations;

  MatrixType m_CachedForceMatrix;

  unsigned int m_NumberOfThreads;

  itk::SimpleFastMutexLock m_Mutex;

  unsigned int m_ElementCounter;

  VectorType _MT_u;
  VectorType _MT_Ku;

  static ITK_THREAD_RETURN_TYPE _productKuNoDiagonalThread(void* arg);
  static ITK_THREAD_RETURN_TYPE _productKuThread(void* arg);

  DeformationFieldPointer m_Deformation;
  DeformationFieldPointer m_InverseDeformation;

  FloatImagePointer _MT_w;
  DeformationFieldPointer _MT_def;

  static ITK_THREAD_RETURN_TYPE _fillDeformationThread(void* arg);

};

#endif
