
#include "TumorEdemaInfiltrationGenerator.h"

#include "itkBSplineInterpolateImageFunction.h"
#include "itkDiscreteGaussianImageFilter.h"
#include "itkImageRegionIterator.h"
#include "itkImageRegionIteratorWithIndex.h"
#include "itkLinearInterpolateImageFunction.h"
#include "itkSignedDanielssonDistanceMapImageFilter.h"
#include "itkWarpImageFilter.h"

#include "itkNearestNeighborInterpolateImageFunction.h"
#include "itkIdentityTransform.h"
#include "itkShrinkImageFilter.h"

#include "vtkIdList.h"
#include "vtkPoints.h"
#include "vtkPointLocator.h"
#include "vtkSmartPointer.h"

#include "createMesh3D.h"

#include "vnl/vnl_math.h"

#include "DTImageInterpolator.h"
#include "DTImageWarpFilter.h"

#include "muException.h"
#include "DynArray.h"
#include "Log.h"

#include "SubsampledImageRegionIterator.h"
#include "VectorBlurImageFilter.h"

#define SOLVER_USE_AVERAGING 0


TumorEdemaInfiltrationGenerator
::TumorEdemaInfiltrationGenerator()
{
  m_InfiltrationIterations = 10;
  m_InfiltrationTimeStep = 0.1;

  m_Lambda = 5.0;
  m_LambdaDamping = 0.95;
  m_LambdaCurrent = 5.0;

  m_ReactionCoefficient = 0.1;

  m_EarlyInfiltrationTime = 1.0;

  m_InfiltrationSolverIterations = 5;

  m_WhiteMatterTensorMultiplier = 10.0;
  m_GrayMatterTensorMultiplier = 1.0;
}

TumorEdemaInfiltrationGenerator
::~TumorEdemaInfiltrationGenerator()
{

}

void
TumorEdemaInfiltrationGenerator
::ReinitializeMesh()
{
  if (m_InitialVTKMesh.GetPointer() != 0)
    m_InitialMesh->SetVTKMesh(m_InitialVTKMesh);
}

void
TumorEdemaInfiltrationGenerator
::ModifyInitialMesh()
{

#if 0

  // Only sets up the mappings, don't really modify the mesh
  // Assumes mesh from mass effect is OK and includes centroids
  //std::cout << "Setting up mesh natural functions, no tumor insertion..." << std::endl;
  delete m_InitialMesh;
  m_InitialMesh = new LinearTetrahedralMesh();
  m_InitialMesh->SetVTKMesh(m_InitialVTKMesh);

#else

  // Insert internal tumor points to help improve infiltration estimation
  //std::cout << "Setting up mesh natural functions and inserting internal tumor points..." << std::endl;

  unsigned int numPoints = m_InitialVTKMesh->GetNumberOfPoints();

  FloatImageSizeType size = m_LabelImage->GetLargestPossibleRegion().GetSize();
  FloatImageSpacingType spacing = m_LabelImage->GetSpacing();

  vtkSmartPointer<vtkPoints> modifiedPoints = vtkSmartPointer<vtkPoints>::New();
  modifiedPoints->Allocate(numPoints);

  double bounds[6];
  for (int dim = 0; dim < 3; dim++)
  { 
    bounds[2*dim] = 0.0;
    bounds[2*dim + 1] = size[dim]*spacing[dim];
  }   
    
  vtkSmartPointer<vtkPointLocator> pLoc =
    vtkSmartPointer<vtkPointLocator>::New();
  pLoc->SetTolerance(m_PointTolerance);
  pLoc->InitPointInsertion(modifiedPoints, bounds, numPoints);

  // Insert original mesh points
  m_InternalInfiltrationPointIds.Clear();

  double x[3];

  for (unsigned int i = 0; i < numPoints; i++)
  {
    m_InitialVTKMesh->GetPoint(i, x);

    vtkIdType id = -1;

    pLoc->InsertUniquePoint(x, id);
  }

  // Insert internal tumor points
  typedef itk::ImageRegionIteratorWithIndex<ByteImageType>
    LabelIteratorType;

  LabelIteratorType labelIt(
    m_LabelImage, m_LabelImage->GetLargestPossibleRegion());
  
  for (labelIt.GoToBegin(); !labelIt.IsAtEnd(); ++labelIt)
  { 
    unsigned char c = labelIt.Get();
    if (c != 5)
      continue;

    ByteImageIndexType ind = labelIt.GetIndex();

    bool ongrid = true;
    for (unsigned int dim = 0; dim < 3; dim++)
      if ((ind[dim] % 2) != 0)
      {
        ongrid = false;
        break;
      }

    if (!ongrid)
      continue;

    ByteImagePointType p;
    m_LabelImage->TransformIndexToPhysicalPoint(ind, p);

    for (unsigned int d = 0; d < 3; d++)
      x[d] = p[d];

    vtkIdType id = -1;

    pLoc->InsertUniquePoint(x, id);

    if (id >= 0)
      m_InternalInfiltrationPointIds.Append(id);
  }

  // Insert centroids to make sure we fill values inside tumor
  for (unsigned int i = 0; i < m_TumorCentroids.GetSize(); i++)
  {
    VectorType c = m_TumorCentroids[i];

    for (unsigned int d = 0; d < 3; d++)
      x[d] = c[d];

    vtkIdType id = -1;
    pLoc->InsertUniquePoint(x, id);

    if (id >= 0)
      m_InternalInfiltrationPointIds.Append(id);
  }

  delete m_InitialMesh;
  m_InitialMesh = new LinearTetrahedralMesh();
  m_InitialMesh->SetVTKMesh(
    createMesh3D(pLoc->GetPoints(), m_BrainDistanceImage, m_UseQHull) );

#endif
}

// Computation of body forces
#if 0

// Compute forces using image gradient
void
TumorEdemaInfiltrationGenerator
::ComputeForces(MatrixType& solutions)
{
  muLogMacro(<< "    Computing body forces with lambda = "
    << m_LambdaCurrent*1e+3 << "\n");
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();
  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();

  FloatImagePointer infilImg = this->GetCurrentInfiltration(2.0);

  typedef itk::BSplineInterpolateImageFunction<FloatImageType, double>
    InterpolatorType;

  InterpolatorType::Pointer infilInterp = InterpolatorType::New();
  infilInterp->SetInputImage(infilImg);
  infilInterp->SetSplineOrder(3);

  solutions.set_size(numPoints, 3);
  solutions.fill(0.0);

  double x[3];

  FloatImagePointType p;

  for (unsigned int i = 0; i < numPoints; i++)
  {
    m_CurrentMesh->GetVTKMesh()->GetPoint(i, x);

    for (int dim = 0; dim < 3; dim++)
      p[dim] = x[dim];

    if (!infilInterp->IsInsideBuffer(p))
      continue;

    // Obtain gradient
    InterpolatorType::CovariantVectorType v =
      infilInterp->EvaluateDerivative(p);

    for (int dim = 0; dim < 3; dim++)
      solutions(i, dim) -= m_LambdaCurrent * v[dim];
  }

  // Reduce lambda by a damping factor after each iteration
  m_LambdaCurrent *= m_LambdaDamping;
}
#else

// Compute forces using FEM derivative
void
TumorEdemaInfiltrationGenerator
::ComputeForces(MatrixType& solutions)
{
  muLogMacro(<< "    Computing body forces with lambda = "
    << m_LambdaCurrent*1e+3 << "\n");
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();
  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();

  solutions.set_size(numPoints, 3);
  solutions.fill(0.0);

  double pt[3];

  for (unsigned int el = 0; el < numElements; el++)
  {
    vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();

    m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, ptIds);

    if (ptIds->GetNumberOfIds() != 4)
      continue;

    MatrixType delN =
      m_CurrentMesh->ComputeShapeFunctionDerivatives(el);

    // Force at centroid
    VectorType f(3, 0.0);
    for (int i = 0; i < 4; i++)
    {
      int global_i = ptIds->GetId(i);
      VectorType dN_i = delN.get_row(i);

      float v = m_InfiltrationNodeSolutions[global_i];

      for (int dim = 0; dim < 3; dim++)
        f[dim] += m_LambdaCurrent * v * dN_i[dim];
    }

    // Distribute to points
    f /= 4;
    for (int i = 0; i < 4; i++)
    {
      int global_i = ptIds->GetId(i);
      for (int dim = 0; dim < 3; dim++)
        solutions(global_i, dim) -= f[dim];
    }

  } // for each el

  // Reduce lambda by a damping factor after each iteration
  m_LambdaCurrent *= m_LambdaDamping;

}
#endif

TumorEdemaInfiltrationGenerator::VectorType
TumorEdemaInfiltrationGenerator
::ProductMass(const VectorType& c)
{
// Returns M * c

  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

  if (c.size() != numPoints)
    muExceptionMacro(<< "ProductMass: invalid input");

  VectorType Mc(numPoints, 0.0);

  for (unsigned int el = 0; el < numElements; el++)
  {
    vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();

    m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, ptIds);

    if (ptIds->GetNumberOfIds() != 4)
      continue;

    float vol = m_CurrentMesh->ComputeElementVolume(el);

    float m_ij_el = vol / 16.0;

    for (unsigned int i = 0; i < 4; i++)
    {
      unsigned int global_i = ptIds->GetId(i);

      for (unsigned int j = 0; j < 4; j++)
      {
        unsigned int global_j = ptIds->GetId(j);

        Mc[global_i] += m_ij_el * c[global_j];
      }
    }

  } // for M*c
/*
  for (unsigned int i = 0; i < numPoints; i++)
    Mc[i] = m_LumpedNodeMasses[i] * c[i];
*/

  return Mc;
}

ITK_THREAD_RETURN_TYPE
TumorEdemaInfiltrationGenerator
::_productDiffusionThread(void* arg)
{
  typedef itk::MultiThreader::ThreadInfoStruct  ThreadInfoType;
  ThreadInfoType * infoStruct = static_cast<ThreadInfoType*>( arg );
  TumorEdemaInfiltrationGenerator* obj = static_cast<TumorEdemaInfiltrationGenerator*>(infoStruct->UserData);

  unsigned int numPoints = obj->m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = obj->m_CurrentMesh->GetNumberOfElements();

  // Update values
  while (true)
  {
    unsigned int el = obj->GetNextElement();

    if (el >= numElements)
      break;

    vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();

    (obj->m_Mutex).Lock();
    obj->m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, ptIds);
    (obj->m_Mutex).Unlock();

    if (ptIds->GetNumberOfIds() != 4)
      continue;

    MatrixType delN =
      obj->m_CurrentMesh->ComputeShapeFunctionDerivatives(el);

    float vol = obj->m_CurrentMesh->ComputeElementVolume(el);

    for (unsigned int i = 0; i < 4; i++)
    {
      unsigned int global_i = ptIds->GetId(i);
      VectorType dN_i = delN.get_row(i);

      VectorType v = obj->m_ElementTensors[el] * dN_i;

      for (unsigned int j = 0; j < 4; j++)
      {
        unsigned int global_j = ptIds->GetId(j);
        VectorType dN_j = delN.get_row(j);

        // Clatz' paper
        //VectorType v = obj->m_ElementTensors[el] * dN_i;
        float d_ij_el = dot_product(dN_j, v) * vol;

        // Hughes' book
        //VectorType v = obj->m_ElementTensors[el] * dN_j;
        //float d_ij_el = dot_product(dN_i, v) * vol;

        (obj->m_Mutex).Lock();
        obj->_MT_Dc[global_i] +=  d_ij_el * obj->_MT_c[global_j];
        (obj->m_Mutex).Unlock();
      }
    }

  } // for D*c

  return ITK_THREAD_RETURN_VALUE;
}

ITK_THREAD_RETURN_TYPE
TumorEdemaInfiltrationGenerator
::_productDiffusionNoDiagonalThread(void* arg)
{
  typedef itk::MultiThreader::ThreadInfoStruct  ThreadInfoType;
  ThreadInfoType * infoStruct = static_cast<ThreadInfoType*>( arg );
  TumorEdemaInfiltrationGenerator* obj = static_cast<TumorEdemaInfiltrationGenerator*>(infoStruct->UserData);

  unsigned int numPoints = obj->m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = obj->m_CurrentMesh->GetNumberOfElements();

  // Update values
  while (true)
  {
    unsigned int el = obj->GetNextElement();

    if (el >= numElements)
      break;

    vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();

    (obj->m_Mutex).Lock();
    obj->m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, ptIds);
    (obj->m_Mutex).Unlock();

    if (ptIds->GetNumberOfIds() != 4)
      continue;

    MatrixType delN =
      obj->m_CurrentMesh->ComputeShapeFunctionDerivatives(el);

    float vol = obj->m_CurrentMesh->ComputeElementVolume(el);

    for (unsigned int i = 0; i < 4; i++)
    {
      unsigned int global_i = ptIds->GetId(i);
      VectorType dN_i = delN.get_row(i);

      VectorType v = obj->m_ElementTensors[el] * dN_i;

      for (unsigned int j = 0; j < 4; j++)
      {
        if (i == j)
          continue;

        unsigned int global_j = ptIds->GetId(j);
        VectorType dN_j = delN.get_row(j);

        if (global_i == global_j)
          continue;

        // Clatz' paper
        //VectorType v = obj->m_ElementTensors[el] * dN_i;
        float d_ij_el = dot_product(dN_j, v) * vol;

        // Hughes' book
        //VectorType v = obj->m_ElementTensors[el] * dN_j;
        //float d_ij_el = dot_product(dN_i, v) * vol;

        (obj->m_Mutex).Lock();
        obj->_MT_Dc[global_i] +=  d_ij_el * obj->_MT_c[global_j];
        (obj->m_Mutex).Unlock();
      }
    }

  } // for D*c

  return ITK_THREAD_RETURN_VALUE;
}

// Returns D*c
TumorEdemaInfiltrationGenerator::VectorType
TumorEdemaInfiltrationGenerator
::ProductDiffusion(const VectorType& c)
{

  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

  if (c.size() != numPoints)
    muExceptionMacro(<< "ProductDiffusion: invalid input size");

  this->_MT_c = c;
  this->_MT_Dc = VectorType(numPoints, 0.0);

  // Try to use half of the number of CPUs, default is min(#cpus, 8)
  itk::MultiThreader::Pointer threader = itk::MultiThreader::New();

  int numThreads = this->GetNumberOfThreads();
  if (numThreads == 0)
    numThreads = threader->GetGlobalDefaultNumberOfThreads() / 4 * 3;
  if (numThreads < 2)
    numThreads = 2;

  this->ResetElementCounter();

  threader->SetNumberOfThreads(numThreads);
  threader->SetSingleMethod(
    &TumorEdemaInfiltrationGenerator::_productDiffusionThread, (void*)this);
  threader->SingleMethodExecute();

  return this->_MT_Dc;

}

// Returns D*c
TumorEdemaInfiltrationGenerator::VectorType
TumorEdemaInfiltrationGenerator
::ProductDiffusionNoDiagonal(const VectorType& c)
{

  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

  if (c.size() != numPoints)
    muExceptionMacro(<< "ProductDiffusionNoDiagonal: invalid input size");

  this->_MT_c = c;
  this->_MT_Dc = VectorType(numPoints, 0.0);

  // Try to use half of the number of CPUs, default is min(#cpus, 8)
  itk::MultiThreader::Pointer threader = itk::MultiThreader::New();

  int numThreads = this->GetNumberOfThreads();
  if (numThreads == 0)
    numThreads = threader->GetGlobalDefaultNumberOfThreads() / 4 * 3;
  if (numThreads < 2)
    numThreads = 2;

  this->ResetElementCounter();

  threader->SetNumberOfThreads(numThreads);
  threader->SetSingleMethod(
    &TumorEdemaInfiltrationGenerator::_productDiffusionNoDiagonalThread, (void*)this);
  threader->SingleMethodExecute();

  return this->_MT_Dc;

}

void
TumorEdemaInfiltrationGenerator
::ComputeMeshVariables()
{
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();
  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();

  m_ElementTensors.Clear();
  m_ElementTensors.Allocate(numElements);
  for (unsigned int el = 0; el < numElements; el++)
    m_ElementTensors.Append(this->GetElementTensor(el));

  // Compute lumped node masses
  m_LumpedNodeMasses.Clear();
  m_LumpedNodeMasses.Initialize(numPoints, 1e-10);

  for (unsigned int el = 0; el < numElements; el++)
  {
    vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();
    m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, ptIds);

    float m = m_CurrentMesh->ComputeElementVolume(el) / 4.0;

    for (unsigned int i = 0; i < ptIds->GetNumberOfIds(); i++)
      m_LumpedNodeMasses[ptIds->GetId(i)] += m;
  }

  // Compute diagonal component of D matrix
  m_DiiList.Clear();
  m_DiiList.Allocate(numPoints);
  for (unsigned int i = 0; i < numPoints; i++)
  {
    float d_ii = 0;

    vtkSmartPointer<vtkIdList> cellIds = vtkSmartPointer<vtkIdList>::New();

    m_CurrentMesh->GetVTKMesh()->GetPointCells(i, cellIds);

    for (unsigned int j = 0; j < cellIds->GetNumberOfIds(); j++)
    {
      unsigned int el = cellIds->GetId(j);

      vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();

      m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, ptIds);

      MatrixType delN = m_CurrentMesh->ComputeShapeFunctionDerivatives(el);

      float vol = m_CurrentMesh->ComputeElementVolume(el);

      vtkIdType local_i = 0;
      for (unsigned int k = 0; k < 4; k++)
        if (ptIds->GetId(k) == i)
        {
          local_i = k;
           break;
        }

      VectorType dN = delN.get_row(local_i);
      VectorType v = m_ElementTensors[el] * dN;

      d_ii += vol * dot_product(dN, v);
    }

    m_DiiList.Append(d_ii + 1e-10);
  }
}

TumorEdemaInfiltrationGenerator::VectorType
TumorEdemaInfiltrationGenerator
::JacobiDiffusion(const VectorType& c)
{
// Returns inv (1 + inv(M)*D*dt) * c
// inv (K) * c where K is
// Kii = 1 + d_ii / m_i * dt
// Kij = d_ij / m_i  * dt

// If RHS is premultiplied by M :
// Returns inv(M + dt*D) * c where M is the lumped mass diag matrix
// // Returns inv(1.0 + dt*D) * c

  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

  if (c.size() != numPoints)
    muExceptionMacro(<< "JacobiDiffusion: invalid input");

  VectorType sol = c;

  for (unsigned int iter = 1; iter <= m_InfiltrationSolverIterations; iter++)
  {
    // Preparation: c+ = c
    VectorType nextSol = c;

    // c+j -= Kij*c
    VectorType vec = this->ProductDiffusionNoDiagonal(sol);
    for (unsigned int i = 0; i < numPoints; i++)
      nextSol[i] -= vec[i] * m_InfiltrationTimeStep / m_LumpedNodeMasses[i];

/*
    for (unsigned int el = 0; el < numElements; el++)
    {
      vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();

      m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, ptIds);

      if (ptIds->GetNumberOfIds() != 4)
        continue;

      MatrixType delN = m_CurrentMesh->ComputeShapeFunctionDerivatives(el);

      float vol = m_CurrentMesh->ComputeElementVolume(el);

      for (unsigned int i = 0; i < 4; i++)
      {
        unsigned int global_i = ptIds->GetId(i);
        VectorType dN_i = delN.get_row(i);

        VectorType v = m_ElementTensors[el] * dN_i;

        for (unsigned int j = 0; j < 4; j++)
        {
          if (i == j)
            continue;

          unsigned int global_j = ptIds->GetId(j);
          VectorType dN_j = delN.get_row(j);

          if (global_i == global_j)
            continue;

          // Clatz' paper
          float d_ij_el = dot_product(dN_j, v) * vol;

          // Hughes' book
          //VectorType v = m_ElementTensors[el] * dN_j;
          //float d_ij_el = dot_product(dN_i, v) * vol;

          nextSol[global_i] -= m_InfiltrationTimeStep * d_ij_el * sol[global_j];
        }
      }
    } // for Dij
*/

    // Divide by Kii
    float maxSolDiff = 0;
    for (unsigned int i = 0; i < numPoints; i++)
    {
      float d_ii = m_DiiList[i];
      float k_ii =
        1.0 + m_DiiList[i] / m_LumpedNodeMasses[i] * m_InfiltrationTimeStep;

      float v = nextSol[i] / k_ii;

      float diff = fabs(v - sol[i]);
      if (diff > maxSolDiff)
        maxSolDiff = diff;

      nextSol[i] = v;
    } //for Dii

#if SOLVER_USE_AVERAGING
    // Average steps, if iters > 1
    if (iter > 1)
    {
      for (unsigned int i = 0; i < numPoints; i++)
        sol[i] = 0.5 * sol[i] + 0.5 * nextSol[i];
    }
    else
    {
      sol = nextSol;
    }
#else
    sol = nextSol;
#endif

    if (maxSolDiff < 1e-8)
      break;

  } // for iter

  // Clamp infiltration values
  for (unsigned int i = 0; i < numPoints; i++)
  {
    float v = sol[i];
    if (vnl_math_isnan(v))
      v = 0.0;
    if (vnl_math_isinf(v))
      v = 0.0;
    if (v < 0.0)
      v = 0.0;
    if (v > 1.0)
      v = 1.0;
    sol[i] = v;
  }

  // Make sure internal infiltration is always 1
  for (unsigned int k = 0; k < m_InternalInfiltrationPointIds.GetSize(); k++)
  {
    unsigned int id = m_InternalInfiltrationPointIds[k];
    sol[id] = 1.0;
  }

  // return sol;

  // Smooth the solution
  VectorType smoothSol(numPoints, 0.0);
  VectorType weights(numPoints, 1e-10);

  for (unsigned int el = 0; el < numElements; el++)
  {
    float vol = m_CurrentMesh->ComputeElementVolume(el);

    vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();

    m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, ptIds);

    if (ptIds->GetNumberOfIds() != 4)
      continue;

    float v = 0.0;

    for (unsigned int i = 0; i < 4; i++)
    {
      unsigned int global_i = ptIds->GetId(i);
      v += sol[global_i];
    }

    v /= 4.0;
    v *= vol;

    for (unsigned int i = 0; i < 4; i++)
    {
      unsigned int global_i = ptIds->GetId(i);
      weights[global_i] += vol;
      smoothSol[global_i] += v;
    }
  } 

  for (unsigned int i = 0; i < numPoints; i++)
    smoothSol[i] /= weights[i];

  return smoothSol;

}

TumorEdemaInfiltrationGenerator::VectorType
TumorEdemaInfiltrationGenerator
::ExplicitDiffusion(const VectorType& c)
{
// Returns c - D*c*dt

  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

  if (c.size() != numPoints)
    muExceptionMacro(<< "ExplicitDiffusion: invalid input");

  VectorType dvec = this->ProductDiffusion(c);

  for (unsigned int i = 0; i < numPoints; i++)
  {
    float v = c[i] - dvec[i] / m_LumpedNodeMasses[i] * m_InfiltrationTimeStep;
    //float v = c[i] - dvec[i] * m_InfiltrationTimeStep;
    if (v < 0.0)
      v = 0.0;
    if (v > 1.0)
      v = 1.0;
    dvec[i] = v;
  }

  return dvec;
}

TumorEdemaInfiltrationGenerator::VectorType
TumorEdemaInfiltrationGenerator
::ImplicitDiffusion(const VectorType& c)
{
// Returns (1 + dt*inv(M)*D) * c
// If RHS premult by M:
// Returns (M + dt*D) * c

  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

  if (c.size() != numPoints)
    muExceptionMacro(<< "ImplicitDiffusion: invalid input");

  VectorType dvec = this->ProductDiffusion(c);

  // c + Dc*m_InfiltrationTimeStep;
  for (unsigned int i = 0; i < numPoints; i++)
  {
    float v = c[i] + dvec[i] / m_LumpedNodeMasses[i] * m_InfiltrationTimeStep;
    //float v = c[i] + dvec[i] * m_InfiltrationTimeStep;
    if (v < 0.0)
      v = 0.0;
    if (v > 1.0)
      v = 1.0;
    dvec[i] = v;
  }

  return dvec;

  //return this->ProductMass(c) + Dc*m_InfiltrationTimeStep;

}

TumorEdemaInfiltrationGenerator::VectorType
TumorEdemaInfiltrationGenerator
::GMRESDiffusion(
  const VectorType& init, const VectorType& rhs, unsigned int maxiters)
{
  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

  VectorType r;
  if (init.magnitude() >=  1e-10)
    //r = rhs - this->ProductDiffusion(init);
    r = rhs - this->ImplicitDiffusion(init);
  else
    r = rhs;

  MatrixType H(maxiters+2, maxiters+2, 0.0);

  MatrixType V(numPoints, maxiters+2, 0.0);

  float rho = r.magnitude();
  V.set_column(1, r / rho);

  float beta = rho;

  float errTol = rhs.magnitude() * 1e-4;

  VectorType y;

  unsigned int k = 0;

  while ((rho > errTol) && (k < maxiters))
  {
    ++k;

    // No precond
    //VectorType Kv = this->ProductDiffusion(V.get_column(k));
    VectorType Kv = this->ImplicitDiffusion(V.get_column(k));
    V.set_column(k+1, Kv);

    for (unsigned int j = 1; j <= k; j++)
    {
      H(j, k) = dot_product(V.get_column(k+1), V.get_column(j));
      V.set_column(k+1, V.get_column(k+1) - V.get_column(j)*H(j, k));
    }
    H(k+1, k) = V.get_column(k+1).magnitude();

    if (H(k+1, k) != 0.0)
      V.set_column(k+1, V.get_column(k+1)/H(k+1, k));

    MatrixType Hk = H.extract(k+1, k, 1, 1);

    VectorType be(k+1, 0.0);
    be[0] = beta;

    MatrixQRType qr(Hk);
    y = qr.solve(be);
    //MatrixSVDType svd(Hk);
    //y = svd.solve(be);

    VectorType errvec = be - Hk*y;

    rho = errvec.squared_magnitude();
//std::cout << "Error iter " << k << " = " << err << std::endl;
  }

  return init + V.extract(numPoints, k, 0, 1) * y;

}

TumorEdemaInfiltrationGenerator::VectorType
TumorEdemaInfiltrationGenerator
::GMRESMass(
  const VectorType& init, const VectorType& rhs, unsigned int maxiters)
{
  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

  VectorType r;
  if (init.magnitude() != 0.0)
    r = rhs - this->ProductMass(init);
  else
    r = rhs;

  MatrixType H(maxiters+2, maxiters+2, 0.0);

  MatrixType V(numPoints, maxiters+2, 0.0);

  float rho = r.magnitude();
  V.set_column(1, r / rho);

  float beta = rho;

  float errTol = rhs.magnitude() * 1e-4;

  VectorType y;

  unsigned int k = 0;

  while ((rho > errTol) && (k < maxiters))
  {
    ++k;

    // No precond
    VectorType Kv = this->ProductMass(V.get_column(k));
    V.set_column(k+1, Kv);

    for (unsigned int j = 1; j <= k; j++)
    {
      H(j, k) = dot_product(V.get_column(k+1), V.get_column(j));
      V.set_column(k+1, V.get_column(k+1) - V.get_column(j)*H(j, k));
    }
    H(k+1, k) = V.get_column(k+1).magnitude();

    if (H(k+1, k) != 0.0)
      V.set_column(k+1, V.get_column(k+1)/H(k+1, k));

    MatrixType Hk = H.extract(k+1, k, 1, 1);

    VectorType be(k+1, 0.0);
    be[0] = beta;

    MatrixQRType qr(Hk);
    y = qr.solve(be);
    //MatrixSVDType svd(Hk);
    //y = svd.solve(be);

    VectorType errvec = be - Hk*y;

    rho = errvec.squared_magnitude();
//std::cout << "Error iter " << k << " = " << rho << std::endl;
  }

  return init + V.extract(numPoints, k, 0, 1) * y;

}

void
TumorEdemaInfiltrationGenerator
::IterateRD()
{
  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

  //VectorType nextInfilSol = m_InfiltrationNodeSolutions;

/*
// Initial diffusion solution, roughly c + inv(M)*dt*D*c
  VectorType nextInfilSol = this->ProductDiffusion(m_InfiltrationNodeSolutions);
  for (unsigned int i = 0; i < numPoints; i++)
  {
    nextInfilSol[i] /= m_LumpedNodeMasses[i];
  }
*/

  // NOTE: explicit step can be unstable
  //VectorType nextInfilSol = this->ExplicitDiffusion(m_InfiltrationNodeSolutions);

  VectorType nextInfilSol = this->JacobiDiffusion(m_InfiltrationNodeSolutions);
  
  // DEBUG: attempt at preconditioning
  //VectorType rhs = this->ProductMass(m_InfiltrationNodeSolutions);
  //VectorType nextInfilSol = this->JacobiDiffusion(rhs);

  for (unsigned int iter = 0; iter < m_InfiltrationSolverIterations; iter++)
  {
    nextInfilSol =
      this->GMRESDiffusion(nextInfilSol, m_InfiltrationNodeSolutions, 5);
      //this->GMRESDiffusion(nextInfilSol, rhs, 5);
  }

  //rhs.set_size(1);

  // Reaction term:
  // Implicit solution by solving quadratic equation for update
  // y is new, x is old
  // (dt*rho) *y^2 + (1-dt*rho)*y + (-x) = 0.0;
  float A = m_InfiltrationTimeStep*m_ReactionCoefficient;
  float B = 1.0 - m_InfiltrationTimeStep*m_ReactionCoefficient;
  for (unsigned int i = 0; i < numPoints; i++)
  {
    //float infil0 = m_InfiltrationNodeSolutions[i];
    float infil0 = nextInfilSol[i];
    nextInfilSol[i] +=
      m_InfiltrationTimeStep * m_ReactionCoefficient * infil0 * (1.0 - infil0);
  }

  // Clamp values
  float sumInfil = 0.0;
  for (unsigned int i = 0; i < numPoints; i++)
  {
    float infil = nextInfilSol[i];
    if (vnl_math_isnan(infil))
      infil = 0.0;
    if (vnl_math_isinf(infil))
      infil = 0.0;
    if (infil < 0.0)
      infil = 0.0;
    if (infil > 1.0)
      infil = 1.0;
    sumInfil += infil;
    nextInfilSol[i] = infil;
  }

  // HACK: Make sure infiltration inside tumor node is always 1
  for (unsigned int k = 0; k < m_InternalInfiltrationPointIds.GetSize(); k++)
  {
    unsigned int id = m_InternalInfiltrationPointIds[k];
    nextInfilSol[id] = 1.0;
  }

  //muLogMacro(<< "      sum(infil) = " << sumInfil << "\n");

  m_InfiltrationNodeSolutions = nextInfilSol;
}

TumorEdemaInfiltrationGenerator::MatrixType
TumorEdemaInfiltrationGenerator
::GetElementTensor(unsigned int el)
{
  // Set up DT interpolator
  DTImageInterpolator dtInterp;
  dtInterp.SetInputImage(m_WarpedDTImage);

  typedef itk::NearestNeighborInterpolateImageFunction<ByteImageType, double>
    NNInterpolatorType;
  NNInterpolatorType::Pointer labelInterp = NNInterpolatorType::New();
  labelInterp->SetInputImage(m_WarpedLabelImage);

  vtkSmartPointer<vtkIdList> elPtIds = vtkSmartPointer<vtkIdList>::New();

  m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, elPtIds);

  if (elPtIds->GetNumberOfIds() != 4)
    return MatrixType(3, 3, 0.0);

  // Compute tetrahedra center
  VectorType x(3, 0.0);
  for (int j = 0; j < 4; j++)
  {
    double pt[3];
    m_CurrentMesh->GetVTKMesh()->GetPoint(elPtIds->GetId(j), pt);

    for (int k = 0; k < 3; k++)
      x[k] += pt[k];
  }
  x /= 4.0;

  DTImagePointType p;
  for (int dim = 0; dim < 3; dim++)
    p[dim] = x[dim];

  DiffusionTensor D = dtInterp.Evaluate(p);

  // Only interested in direction, magnitude is set based on class
  // NOTE: This will disregard DTI destruction simulation
  //D.ForceUnitTrace();

  DiffusionTensor::MatrixType M = D.GetMatrix();

  float trM = 0.0;
  for (int k = 0; k < 3; k++)
    trM += M(k, k);

  float md = trM / 3.0;

  // Scalar multiplication based on label majority within cell
  DynArray<unsigned int> countLabels;
  countLabels.Initialize(10, 0);
  for (int j = 0; j < 4; j++)
  {
    double pt[3];
    m_CurrentMesh->GetVTKMesh()->GetPoint(elPtIds->GetId(j), pt);

    ByteImagePointType p;
    for (int k = 0; k < 3; k++)
      p[k] = pt[k];

    if (!labelInterp->IsInsideBuffer(p))
    {
      countLabels[0] += 1;
      continue;
    }

    unsigned int label = (unsigned int)labelInterp->Evaluate(p);
    //if (label >= 6)
    //  muExceptionMacro(<< "Unexpected label value: " << label);
    countLabels[label] += 1;
  }

  unsigned int maxLabel = 0;
  unsigned int maxCount = countLabels[0];
  for (unsigned int i = 1; i < countLabels.GetSize(); i++)
  {
    if (countLabels[i] > maxCount)
    {
      maxCount = countLabels[i];
      maxLabel = i;
    }
  }

  // Diffusion tensor modifier for different tissues
  float alpha = m_WhiteMatterTensorMultiplier;
  float beta = m_GrayMatterTensorMultiplier;

  if (maxLabel == 1)
  {
    M *= alpha;
  }
  else if (maxLabel == 2)
  {
    //M *= beta;

    // Disregard directional component in gm?
    M.set_identity();
    M *= md * beta;
  }
  else if (maxLabel == 4)
  {
    // No diffusion in falx
    M.fill(0.0);
  }
  else if (maxLabel == 5)
  {
    // Assume tensors inside tumor already has roughly max MD, normalized to 1
    M.set_identity();
    M *= alpha;
  }
  else
  {
    M.set_identity();
    M *= md * beta * 0.01;

    //M.fill(0.0);
  }

  // Add minimum diffusivity in cardinal directions
  for (int dim = 0; dim < 3; dim++)
    M(dim, dim) += 1e-4;

  return M;

}

TumorEdemaInfiltrationGenerator::FloatImagePointer
TumorEdemaInfiltrationGenerator
::ComputeInfiltration()
{
  // Check inputs
  if (m_LabelImage.IsNull())
    muExceptionMacro(<< "No label image specified");

  if (m_DTImage.IsNull())
    muExceptionMacro(<< "No DT image specified");

  if (m_InitialVTKMesh.GetPointer() == 0)
    muExceptionMacro(<< "Initial mesh not specified");

  m_WarpedLabelImage = m_LabelImage;
  m_WarpedDTImage = m_DTImage;

  this->ComputeTumorCentroids();

  // Set up the mappings
  this->ModifyInitialMesh();

  delete m_CurrentMesh;
  m_CurrentMesh = new LinearTetrahedralMesh();
  m_CurrentMesh->SetVTKMesh(m_InitialMesh->GetVTKMesh());

  this->DetermineBCs();

  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

std::cout << "Start infil simulation with " << numElements << " elements and " << numPoints << " points" << std::endl;

  // Initialize compounded deformation solution with zeros
  m_NodeCompSolutions = MatrixType(numPoints, 3, 0.0);

  // Allocate deformation field
  DeformationFieldType::Pointer currDefImg = DeformationFieldType::New();
  // NOTE: CopyInformation fails for vector image
  //currDefImg->CopyInformation(m_LabelImage);
  currDefImg->SetDirection(m_LabelImage->GetDirection());
  currDefImg->SetOrigin(m_LabelImage->GetOrigin());
  currDefImg->SetSpacing(m_LabelImage->GetSpacing());
  currDefImg->SetRegions(m_LabelImage->GetLargestPossibleRegion());
  currDefImg->Allocate();

  m_Deformation = DeformationFieldType::New();
  //m_Deformation->CopyInformation(m_LabelImage);
  m_Deformation->SetDirection(m_LabelImage->GetDirection());
  m_Deformation->SetOrigin(m_LabelImage->GetOrigin());
  m_Deformation->SetSpacing(m_LabelImage->GetSpacing());
  m_Deformation->SetRegions(m_LabelImage->GetLargestPossibleRegion());
  m_Deformation->Allocate();

  m_InverseDeformation = DeformationFieldType::New();
  //m_InverseDeformation->CopyInformation(m_LabelImage);
  m_InverseDeformation->SetDirection(m_LabelImage->GetDirection());
  m_InverseDeformation->SetOrigin(m_LabelImage->GetOrigin());
  m_InverseDeformation->SetSpacing(m_LabelImage->GetSpacing());
  m_InverseDeformation->SetRegions(m_LabelImage->GetLargestPossibleRegion());
  m_InverseDeformation->Allocate();

  DisplacementType zeroDisp;
  zeroDisp.Fill(0.0);

  m_Deformation->FillBuffer(zeroDisp);
  m_InverseDeformation->FillBuffer(zeroDisp);

  unsigned int earlyIters = (unsigned int)(m_EarlyInfiltrationTime / m_InfiltrationTimeStep);

  unsigned int interleavedDefIters = 0;
  if (m_DeformationIterations > 0)
  {
#if 1
    if (m_InfiltrationIterations >= m_DeformationIterations)
      interleavedDefIters = m_InfiltrationIterations / m_DeformationIterations;
    else
      interleavedDefIters = 1;
#else
    if ((m_InfiltrationIterations-earlyIters) >= m_DeformationIterations)
      interleavedDefIters = (m_InfiltrationIterations-earlyIters) / m_DeformationIterations;
    else
      interleavedDefIters = 1;
#endif
    if (interleavedDefIters < 1)
      interleavedDefIters = 1;
  }

  unsigned int countBodyDefIters = 0;

  m_LambdaCurrent = m_Lambda;

  // Define iterators
  typedef itk::ImageRegionIteratorWithIndex<FloatImageType> FloatIteratorType;
  typedef itk::ImageRegionIterator<ByteImageType> ByteIteratorType;
  //typedef itk::ImageRegionIterator<DTImageType> DTIteratorType;

  ByteIteratorType labelIt(m_LabelImage,
    m_LabelImage->GetLargestPossibleRegion());

  // Initialize to field to one within tumor, blur a bit
  {
    // Allocate infiltration
    FloatImagePointer infilImg = FloatImageType::New();
    infilImg->CopyInformation(m_LabelImage);
    infilImg->SetRegions(m_LabelImage->GetLargestPossibleRegion());
    infilImg->Allocate();

    FloatIteratorType infilIt(infilImg, infilImg->GetLargestPossibleRegion());
    labelIt.GoToBegin();
    infilIt.GoToBegin();
    while (!labelIt.IsAtEnd())
    {
      unsigned char label = labelIt.Get();

      if (label == 5)
        infilIt.Set(1.0);
      else
        infilIt.Set(0.0);

      ++labelIt;
      ++infilIt; 
    }

    typedef itk::DiscreteGaussianImageFilter<FloatImageType, FloatImageType>
      BlurType;

    BlurType::Pointer blurf = BlurType::New();
    blurf->SetInput(infilImg);
    blurf->SetVariance(2.0);

    blurf->Update();

    m_InfiltrationImage = blurf->GetOutput();
  }

  // Assign infiltration image values to initial node solutions
  m_InfiltrationNodeSolutions = VectorType(numPoints, 0.0);

  typedef itk::LinearInterpolateImageFunction<FloatImageType, double>
    LinearInterpolatorType;
  LinearInterpolatorType::Pointer infilInterp = LinearInterpolatorType::New();
  infilInterp->SetInputImage(m_InfiltrationImage);

#if 0
  // Initialize using average cell center values
  for (unsigned int el = 0; el < numElements; el++)
  {
    vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();

    m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, ptIds);

    FloatImagePointType centroid;
    centroid.Fill(0.0);

    double pt[3];
    for (unsigned int k = 0; k < 4; k++)
    {
      unsigned int id = ptIds->GetId(k);
      m_CurrentMesh->GetVTKMesh()->GetPoint(id, pt);
      for (int d = 0; d < 3; d++)
        centroid[d] += pt[d];
    }

    for (int d = 0; d < 3; d++)
      centroid[d] /= 4.0;

    if (!infilInterp->IsInsideBuffer(centroid))
      continue;

    float infil = infilInterp->Evaluate(centroid);

    for (unsigned int k = 0; k < 4; k++)
    {
      unsigned int id = ptIds->GetId(k);
      m_InfiltrationNodeSolutions[id] += infil / 4.0;
    }
  }
#else
  // Initialize using point values
  for (unsigned int i = 0; i < numPoints; i++)
  {
    double pt[3];

    m_CurrentMesh->GetVTKMesh()->GetPoint(i, pt);

    FloatImagePointType p;
    for (unsigned int dim = 0; dim < 3; dim++)
      p[dim] = pt[dim];

    if (!infilInterp->IsInsideBuffer(p))
      continue;

    float infil = infilInterp->Evaluate(p);

    m_InfiltrationNodeSolutions[i] = infil;
  }
#endif

  // Make sure initial solutions have max of 1, to speed up est process
  float maxInitInfil = 0.0;
  for (unsigned int k = 0; k < numPoints; k++)
  {
    float v = m_InfiltrationNodeSolutions[k];
    if (v > maxInitInfil)
      maxInitInfil = v;
  }

  if (maxInitInfil == 0.0)
    muExceptionMacro(<< "Initial infiltration values all zero");

  for (unsigned int k = 0; k < numPoints; k++)
  {
    float v = m_InfiltrationNodeSolutions[k] / maxInitInfil;
    if (v < 0.0)
      v = 0.0;
    if (v > 1.0)
      v = 1.0;
    m_InfiltrationNodeSolutions[k] = v;
  }

  // Compute node and element variables for current mesh
  this->ComputeMeshVariables();

  // Iterate reaction-diffusion
  float t = 0.0;
  bool done_early = false;
  for (unsigned int iter = 0; iter < m_InfiltrationIterations; iter++)
  {
    muLogMacro(<< "  * Infiltration update at t = " << t << "\n");

    // Copy current infiltration image as early infiltration
    if (t >= m_EarlyInfiltrationTime && !done_early)
    {
      muLogMacro(<< "  - Storing early infiltration values at t = " << t << "\n");
      m_EarlyInfiltrationImage = this->GetCurrentInfiltration(2.0);

      done_early = true;
    }

    this->IterateRD();

    // Interleaved body force simulations
    if ( (iter % interleavedDefIters == 0) && (countBodyDefIters < m_DeformationIterations) )
    {
      this->ComputeForces(m_CachedForceMatrix);

      this->UpdateSolutionsJacobi();

      this->RefineSolutionsGMRES();

      m_CachedForceMatrix.set_size(1, 1);

      this->UpdateDeformationField(currDefImg);

      this->InvertAccumulateDeformationField(
        m_Deformation, m_InverseDeformation, currDefImg);

      // Update mesh using compounded forward map
      this->UpdateMesh(m_InverseDeformation);

      // Warp dti using m_Deformation
      DTImageWarpFilter::Pointer dtWarper = DTImageWarpFilter::New();
      dtWarper->SetInput(m_DTImage);
      dtWarper->SetDeformationField(m_Deformation);
      dtWarper->SetInverseDeformationField(m_InverseDeformation);
      dtWarper->Update();
      m_WarpedDTImage = dtWarper->GetOutput();

      // Warp label image using m_Deformation
      typedef itk::WarpImageFilter<
        ByteImageType, ByteImageType, DeformationFieldType> ByteWarperType;

      typedef itk::NearestNeighborInterpolateImageFunction<
        ByteImageType, double> NNInterpolatorType;

      ByteWarperType::Pointer warper = ByteWarperType::New();
      warper->SetEdgePaddingValue(0);
      warper->SetInput(m_LabelImage);
      warper->SetInterpolator(NNInterpolatorType::New());
      warper->SetOutputDirection(m_LabelImage->GetDirection());
      warper->SetOutputOrigin(m_LabelImage->GetOrigin());
      warper->SetOutputSpacing(m_LabelImage->GetSpacing());
      //warper->SetOutputSize(m_LabelImage->GetLargestPossibleRegion().GetSize());
#if ITK_VERSION_MAJOR >= 4
      warper->SetDisplacementField(m_Deformation);
#else
      warper->SetDeformationField(m_Deformation);
#endif
      warper->Update();

      m_WarpedLabelImage = warper->GetOutput();


      this->ComputeMeshVariables();

      countBodyDefIters++;
    }

    t += m_InfiltrationTimeStep;
  }

  for (; countBodyDefIters < m_DeformationIterations; countBodyDefIters++)
  {
    this->ComputeForces(m_CachedForceMatrix);

    this->UpdateSolutionsJacobi();

    this->RefineSolutionsGMRES();

    m_CachedForceMatrix.set_size(1, 1);

    this->UpdateDeformationField(currDefImg);

    this->InvertAccumulateDeformationField(
      m_Deformation, m_InverseDeformation, currDefImg);

    // Update mesh using compounded forward map
    this->UpdateMesh(m_InverseDeformation);

    this->ComputeMeshVariables();
  }

  // Fill in infiltration image using interpolated values from mesh
  m_InfiltrationImage = this->GetCurrentInfiltration(2.0);

  delete m_InitialMesh;
  m_InitialMesh = 0;
  m_CurrentMesh->ClearMappings(); // Keep final VTK mesh, but remove mappings

  // Smooth the deformation field
  // Helps the inversion process
  //typedef VectorMeanImageFilter<DeformationFieldType, DeformationFieldType>
  typedef VectorBlurImageFilter<DeformationFieldType, DeformationFieldType>
    DeformationSmootherType;

  // Smoothen final accumulated deformations
  DeformationSmootherType::Pointer defsmoother = DeformationSmootherType::New();
  defsmoother->SetKernelWidth(2.0);
  defsmoother->SetInput(m_InverseDeformation);
  defsmoother->Update();
  m_InverseDeformation = defsmoother->GetOutput();

  DeformationSmootherType::Pointer defsmoother2 = DeformationSmootherType::New();
  defsmoother2->SetKernelWidth(2.0);
  defsmoother2->SetInput(m_Deformation);
  defsmoother2->Update();
  m_Deformation = defsmoother2->GetOutput();

  return m_InfiltrationImage;

}

float
TumorEdemaInfiltrationGenerator
::ComputeInfiltrationSolution(VectorType& x)
{
  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();

  vtkUnstructuredGrid* initialVTKMesh = m_CurrentMesh->GetVTKMesh();

  float sol = 0.0;

  // Find closest points
  LinearTetrahedralMesh::PointType q;
  for (unsigned int dim = 0; dim < 3; dim++)
    q[dim] = x[dim];

  unsigned int numSearchPts = 5;

  LinearTetrahedralMesh::PointIDVectorType closestPoints =
    m_CurrentMesh->FindClosestPoints(q, numSearchPts);

  // Only consider elements that contain the closest point/nodes
  bool foundEl = false;

  for (unsigned int i = 0; i < numSearchPts; i++)
  {
    unsigned int closestPt = closestPoints[i];

    vtkSmartPointer<vtkIdList> cellIds = vtkSmartPointer<vtkIdList>::New();

    initialVTKMesh->GetPointCells(closestPt, cellIds);

    for (unsigned int j = 0; j < cellIds->GetNumberOfIds(); j++)
    {
      unsigned int el = cellIds->GetId(j);

      VectorType N = m_CurrentMesh->ComputeShapeFunctions(el, x);
  
      // Inside/outside test
      float minN = N[0];
      for (unsigned int k = 1; k < 4; k++)
      {
        if (N[k] < minN)
          minN = N[k];
      }
      if (minN < -1e-2)
        continue;
    
      vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();
    
      initialVTKMesh->GetCellPoints(el, ptIds);

      if (ptIds->GetNumberOfIds() != 4)
        continue;

      for (unsigned int k = 0; k < ptIds->GetNumberOfIds(); k++)
        sol += m_InfiltrationNodeSolutions[ptIds->GetId(k)] * N[k];
    
      foundEl = true;
    
      // A global point can only be inside one tetrahedron
      break;
    }

    if (foundEl)
      break;

  } // for i

  if (sol < 0.0)
    sol = 0.0;
  if (sol > 1.0)
    sol = 1.0;

  return sol;
}

ITK_THREAD_RETURN_TYPE
TumorEdemaInfiltrationGenerator
::_fillDiffusionThread(void* arg)
{
  typedef itk::MultiThreader::ThreadInfoStruct  ThreadInfoType;
  ThreadInfoType * infoStruct = static_cast<ThreadInfoType*>( arg );
  TumorEdemaInfiltrationGenerator* obj =
    static_cast<TumorEdemaInfiltrationGenerator*>(infoStruct->UserData);

  unsigned int numPoints = obj->m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = obj->m_CurrentMesh->GetNumberOfElements();

  FloatImageSizeType size =
    obj->m_LabelImage->GetLargestPossibleRegion().GetSize();
  FloatImageSpacingType spacing =
    obj->m_LabelImage->GetSpacing();
  FloatImagePointType orig =
    obj->m_LabelImage->GetOrigin();

  while (true)
  {
    unsigned int el = obj->GetNextElement();

    if (el >= numElements)
      break;

    vtkSmartPointer<vtkIdList> ptIds = vtkSmartPointer<vtkIdList>::New();

    (obj->m_Mutex).Lock();
    obj->m_CurrentMesh->GetVTKMesh()->GetCellPoints(el, ptIds);
    (obj->m_Mutex).Unlock();

    // Compute bounding box of this element
    long minext[3];
    long maxext[3];

    for (int d = 0; d < 3; d++)
    {
      minext[d] = size[d]-1;
      maxext[d] = 0;
    }
  
    for (unsigned int k = 0; k < 4; k++)
    {
      double x[3];
      (obj->m_Mutex).Lock();
      obj->m_CurrentMesh->GetVTKMesh()->GetPoint(ptIds->GetId(k), x);
      (obj->m_Mutex).Unlock();

      FloatImagePointType p;
      for (int d = 0; d < 3; d++)
        p[d] = x[d];

      FloatImageIndexType ind;
      obj->m_LabelImage->TransformPhysicalPointToIndex(p, ind);

      for (int d = 0; d < 3; d++)
      {
        if (ind[d] < minext[d])
          minext[d] = ind[d];
        if (ind[d] > maxext[d])
          maxext[d] = ind[d];
      } 
    }

    // Expand bounding box slightly
    for (int d = 0; d < 3; d++)
    {
      minext[d] -= 5;
      maxext[d] += 5;

      if (minext[d] < 0)
        minext[d] = 0;
      if (maxext[d] >= (long)size[d])
        maxext[d] = size[d]-1;
    }

    // Get the relevant node values
    float nodeValues[4];
    for (unsigned int k = 0; k < 4; k++)
    {
      float infil_k = obj->m_InfiltrationNodeSolutions[ptIds->GetId(k)];
      if (infil_k < 0.0)
        infil_k = 0.0;
      if (infil_k > 1.0)
        infil_k = 1.0;
      nodeValues[k] = infil_k;
    }
    
    // Fill in image values within bounding box
    for (long k = minext[2]; k <= maxext[2]; k++)
      for (long j = minext[1]; j <= maxext[1]; j++)
        for (long i = minext[0]; i <= maxext[0]; i++)
        {
          FloatImageIndexType ind = {{i, j, k}};

          if (obj->m_LabelImage->GetPixel(ind) == 0)
            continue;

          ByteImagePointType p;
          obj->m_LabelImage->TransformIndexToPhysicalPoint(ind, p);

          VectorType x(3);
          for (int d = 0; d < 3; d++)
            x[d] = p[d];

          VectorType N = obj->m_CurrentMesh->ComputeShapeFunctions(el, x);

          float minN = N[0];
          float sumN = N[0];
          for (int t = 1; t < 4; t++)
          {
            if (N[t] < minN)
              minN = N[t];
            sumN += N[t];
          }

          // Inside / outside test
          if (minN < -1e-2)
            continue;

          float sol = 0; 
          for (int t = 0; t < 4; t++)
            sol += N[t] * nodeValues[t];

          // Compound solution values

          (obj->m_Mutex).Lock();
          obj->_MT_weight->SetPixel(ind,
            obj->_MT_weight->GetPixel(ind) + sumN);
          obj->_MT_sol->SetPixel(ind,
            obj->_MT_sol->GetPixel(ind) + sol);
          (obj->m_Mutex).Unlock();

        } // for bbox pixels

  } // for el

  return ITK_THREAD_RETURN_VALUE;
}



// Fill in infiltration image using per-element iterations
TumorEdemaInfiltrationGenerator::FloatImagePointer
TumorEdemaInfiltrationGenerator
::GetCurrentInfiltration(float blurVar)
{
  unsigned int numPoints = m_CurrentMesh->GetNumberOfPoints();
  unsigned int numElements = m_CurrentMesh->GetNumberOfElements();

  FloatImageSizeType size = m_LabelImage->GetLargestPossibleRegion().GetSize();
  FloatImageSpacingType spacing = m_LabelImage->GetSpacing();
  FloatImagePointType orig = m_LabelImage->GetOrigin();

  FloatImagePointer infilImg = FloatImageType::New();
  infilImg->CopyInformation(m_LabelImage);
  infilImg->SetRegions(m_LabelImage->GetLargestPossibleRegion());
  infilImg->Allocate();
  infilImg->FillBuffer(0);

  FloatImagePointer weightImg = FloatImageType::New();
  weightImg->CopyInformation(m_LabelImage);
  weightImg->SetRegions(m_LabelImage->GetLargestPossibleRegion());
  weightImg->Allocate();
  weightImg->FillBuffer(0);

  // Fill in sol and weights
  _MT_sol = infilImg;
  _MT_weight = weightImg;

  itk::MultiThreader::Pointer threader = itk::MultiThreader::New();

  int numThreads = this->GetNumberOfThreads();
  if (numThreads == 0)
    numThreads = threader->GetGlobalDefaultNumberOfThreads() / 4 * 3;
  if (numThreads < 2)
    numThreads = 2;

  this->ResetElementCounter();

  threader->SetNumberOfThreads(numThreads);
  threader->SetSingleMethod(
    &TumorEdemaInfiltrationGenerator::_fillDiffusionThread, (void*)this);
  threader->SingleMethodExecute();

  // Clamp final image values
  typedef itk::ImageRegionIteratorWithIndex<FloatImageType> FloatIteratorType;
  FloatIteratorType it(infilImg, infilImg->GetLargestPossibleRegion());

  float maxV = 0.0;

  for (it.GoToBegin(); !it.IsAtEnd(); ++it)
  {
    FloatImageIndexType ind = it.GetIndex();

    float w = weightImg->GetPixel(ind);

    if (w < 1e-10)
      continue;

    float sol = it.Get() / w;

    if (sol < 0.0)
      sol = 0.0;

    //if (m_WarpedLabelImage->GetPixel(ind) != 5 && sol > maxV)
    if (sol > maxV)
      maxV = sol;

    it.Set(sol);
  }

  if (blurVar > 0)
  {
    // Blur images in case of gaps in tetrahedras
    typedef itk::DiscreteGaussianImageFilter<FloatImageType, FloatImageType>
      BlurType;

    BlurType::Pointer blurf = BlurType::New();
    blurf->SetInput(infilImg);
    blurf->SetVariance(blurVar);
    blurf->Update();

    infilImg = blurf->GetOutput();
  }

  return infilImg;
}
