
////////////////////////////////////////////////////////////////////////////////
//
// Simulate tumor and edema infiltration using DTI-guided reaction-diffusion
// Mass effect is also computed from infiltration gradient
//
////////////////////////////////////////////////////////////////////////////////

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

#ifndef _TumorEdemaInfiltrationGenerator_h
#define _TumorEdemaInfiltrationGenerator_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 "LinearTetrahedralMesh.h"
#include "TumorMassEffectGenerator.h"

#include "DTImage.h"

class TumorEdemaInfiltrationGenerator: public TumorMassEffectGenerator
{

public:

  TumorEdemaInfiltrationGenerator();
  virtual ~TumorEdemaInfiltrationGenerator();

  inline void SetDTImage(DTImagePointer dti)
  { m_DTImage = dti; }

  inline void SetInfiltrationIterations(unsigned int n)
  { m_InfiltrationIterations = n; }

  inline void SetInfiltrationTimeStep(float d)
  { m_InfiltrationTimeStep = d; }

  inline void SetEarlyInfiltrationTime(float t)
  { m_EarlyInfiltrationTime = t; }

  inline void SetLambda(float d)
  { m_Lambda = d; m_LambdaCurrent = d; }
  inline float GetLambda() { return m_Lambda; }

  inline void SetLambdaDamping(float d)
  { m_LambdaDamping = d; }

  // Must be called before computing deformation due to mass effect
  FloatImagePointer ComputeInfiltration();

  inline FloatImagePointer GetEarlyInfiltration()
  { return m_EarlyInfiltrationImage; }

  // Hack:
  // Reinitialize mesh using original, unmodified input VTK mesh
  void ReinitializeMesh();

  inline void SetInfiltrationSolverIterations(unsigned int n)
  { m_InfiltrationSolverIterations = n; }

  inline void SetReactionCoefficient(float d)
  { m_ReactionCoefficient = d; }
  inline void SetWhiteMatterTensorMultiplier(float d)
  { m_WhiteMatterTensorMultiplier = d; }
  inline void SetGrayMatterTensorMultiplier(float d)
  { m_GrayMatterTensorMultiplier = d; }

protected:

  virtual void ModifyInitialMesh();

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

  DiffusionTensor::MatrixType GetElementTensor(unsigned int el);

  // Compute tensor, Dii, and node masses
  void ComputeMeshVariables();

  float ComputeInfiltrationSolution(VectorType& x);

  void IterateRD();

  VectorType ProductMass(const VectorType& c);
  VectorType GMRESMass(
    const VectorType& init, const VectorType& rhs, unsigned int maxiters);

  VectorType ExplicitDiffusion(const VectorType& c);
  VectorType ImplicitDiffusion(const VectorType& c);

  VectorType JacobiDiffusion(const VectorType& init);

  VectorType ProductDiffusion(const VectorType& c);
  VectorType ProductDiffusionNoDiagonal(const VectorType& c);
  VectorType GMRESDiffusion(
    const VectorType& init, const VectorType& rhs, unsigned int maxiters);

//  VectorType ProductReaction(const VectorType& c);
//  void GMRESReaction(
//    const VectorType& init, const VectorType& rhs, unsigned int maxiters);

  FloatImagePointer GetCurrentInfiltration(float blurVar = 0);

private:

  // Coupling factor for mass effect
  float m_Lambda;

  // Damping factor per iteration
  float m_LambdaDamping;

  // Current value for damped lambda
  float m_LambdaCurrent;

  unsigned int m_InfiltrationIterations;
  float m_InfiltrationTimeStep;

  FloatImagePointer m_InfiltrationImage;
  FloatImagePointer m_EarlyInfiltrationImage;

  float m_EarlyInfiltrationTime;

  DTImagePointer m_DTImage;
  DTImagePointer m_WarpedDTImage;

  VectorType m_InfiltrationNodeSolutions;

  DynArray<float> m_LumpedNodeMasses;
  DynArray<float> m_DiiList;

  DynArray<long> m_InternalInfiltrationPointIds;

  float m_ReactionCoefficient;

  float m_WhiteMatterTensorMultiplier;
  float m_GrayMatterTensorMultiplier;

  ByteImagePointer m_WarpedLabelImage;

  DynArray<DiffusionTensor::MatrixType> m_ElementTensors;

  unsigned int m_InfiltrationSolverIterations;

  itk::SimpleFastMutexLock m_Mutex;

  VectorType _MT_c;
  VectorType _MT_Dc;

  static ITK_THREAD_RETURN_TYPE _productDiffusionThread(void* arg);
  static ITK_THREAD_RETURN_TYPE _productDiffusionNoDiagonalThread(void* arg);

  FloatImagePointer _MT_sol;
  FloatImagePointer _MT_weight;

  static ITK_THREAD_RETURN_TYPE _fillDiffusionThread(void* arg);


};

#endif
