
#include "DiffusionTensor.h"

#include "vnl/vnl_math.h"
#include "vnl/vnl_numeric_traits.h"

#include <cfloat>
#include <cmath>
#include <iostream>

DiffusionTensor::
DiffusionTensor()
{
  // Use a "small" isotropic tensor as default
  MatrixType M(3, 3);
  M.set_identity();
  M *= 1e-10;

  this->FromMatrix(M);
}

DiffusionTensor::
DiffusionTensor(const DiffusionTensor& dt)
{
  for (unsigned int i = 0; i < 6; i++)
    m_Elements[i] = dt.m_Elements[i];
}

DiffusionTensor::
~DiffusionTensor()
{

}

const DiffusionTensor&
DiffusionTensor::
operator=(const DiffusionTensor& dt)
{
  for (unsigned int i = 0; i < 6; i++)
    m_Elements[i] = dt.m_Elements[i];
  return *this;
}

const DiffusionTensor&
DiffusionTensor::
operator*=(double d)
{
  for (unsigned int i = 0; i < 6; i++)
    m_Elements[i] = (ScalarType)(m_Elements[i] * d);
  return *this;
}

const DiffusionTensor&
DiffusionTensor::
operator/=(double d)
{
  for (unsigned int i = 0; i < 6; i++)
    m_Elements[i] = (ScalarType)(m_Elements[i] / d);
  return *this;
}

DiffusionTensor
DiffusionTensor::
operator+(const DiffusionTensor& dt)
{
  DiffusionTensor s;
  for (unsigned int i = 0; i < 6; i++)
    s.m_Elements[i] = this->m_Elements[i] + dt.m_Elements[i];
  return s;
}

DiffusionTensor
DiffusionTensor::
operator-(const DiffusionTensor& dt)
{
  DiffusionTensor s;
  for (unsigned int i = 0; i < 6; i++)
    s.m_Elements[i] = this->m_Elements[i] - dt.m_Elements[i];
  return s;
}

DiffusionTensor
DiffusionTensor::
operator*(double d)
{
  DiffusionTensor s;
  for (unsigned int i = 0; i < 6; i++)
    s.m_Elements[i] = (ScalarType)(this->m_Elements[i] * d);
  return s;
}

DiffusionTensor
DiffusionTensor::
operator/(double d)
{
  DiffusionTensor s;
  for (unsigned int i = 0; i < 6; i++)
    s.m_Elements[i] = (ScalarType)(this->m_Elements[i] / d);
  return s;
}

DiffusionTensor::MatrixType
DiffusionTensor::
GetMatrix() const
{
  MatrixType m(3, 3);

  unsigned int k = 0;
  for (unsigned int i = 0; i < 3; i++)
    for (unsigned int j = i; j < 3; j++)
    {
      double t = m_Elements[k++];
      m(i, j) = t;
      m(j, i) = t;
    }

  return m;
}

void
DiffusionTensor::
FromMatrix(const MatrixType& m)
{

#if 0
  if (m.rows() != 3 || m.columns() != 3)
  {
    muExceptionMacro(<< "Invalid matrix for diffusion tensor");
    return;
  }
#endif

  int k = 0;
  for (int i = 0; i < 3; i++)
    for (int j = i; j < 3; j++)
    {
      double d = m(i, j);
      if (vnl_math_isinf(d))
        d = 0.0;
      if (vnl_math_isnan(d))
        d = 0.0;
      m_Elements[k++] = (ScalarType)d;
    }

}

void
DiffusionTensor::
SetElements(const ScalarType* elems)
{
  if (elems == NULL)
  {
    std::cerr << "[DiffusionTensor::SetElements] NULL argument" << std::endl;
    return;
  }

  for (unsigned int i = 0; i < 6; i++)
  {
    ScalarType d = elems[i];
    if (vnl_math_isinf(d))
      d = 0.0;
      //d = vnl_numeric_traits<ScalarType>::maxval;
    if (vnl_math_isnan(d))
      d = 0.0;
    m_Elements[i] = d;
  }
}

void
DiffusionTensor::
SetElementAt(unsigned int i, double v)
{

#if 0
  if (i >= 6)
  {
    std::cerr << "[DiffusionTensor::SetElementAt] Invalid index" << std::endl;
    return;
  }
#endif

  if (vnl_math_isinf(v))
    v = 0.0;
  if (vnl_math_isnan(v))
    v = 0.0;

  m_Elements[i] = (ScalarType)v;
}

void
DiffusionTensor::
SetElementAt(unsigned int i, unsigned int j, double v)
{

  if (i >= 3 || j >= 3)
  {
    std::cerr << "[DiffusionTensor::SetElementAt] Invalid index" << std::endl;
    return;
  }

  unsigned int r = (i < j) ? i : j;
  unsigned int c = (i > j) ? i : j;

  unsigned int k = 0;
  if (r > 0)
    k = r + 1;
  k += c;

  if (vnl_math_isinf(v))
    v = 0.0;
  if (vnl_math_isnan(v))
    v = 0.0;

  m_Elements[k] = (ScalarType)v;
}

void
DiffusionTensor::
ForcePositiveDefinite()
{
  MatrixType M = this->GetMatrix();

  MatrixEigenType eig(M);
  for (unsigned int i = 0; i < 3; i++)
    if (eig.D(i, i) < 1e-10)
      eig.D(i, i) = 1e-10;

  M = eig.recompose();

  this->FromMatrix(M);
}

void
DiffusionTensor::
ForceUnitTrace()
{
  MatrixType M = this->GetMatrix();

  MatrixEigenType eig(M);

  double sumd = FLT_EPSILON;
  for (unsigned int i = 0; i < 3; i++)
    sumd +=  eig.D(i, i);

  for (unsigned int i = 0; i < 3; i++)
    eig.D(i, i) = eig.D(i, i) / sumd;

  M = eig.recompose();

  this->FromMatrix(M);
}

void
DiffusionTensor::
AddMeanDiffusivity(double d)
{
  MatrixType M = this->GetMatrix();

  MatrixEigenType eig(M);

  double sumd = FLT_EPSILON;
  for (unsigned int i = 0; i < 3; i++)
    sumd += eig.D(i, i);

  for (unsigned int i = 0; i < 3; i++)
  {
    double frac = eig.D(i, i) / sumd;
    eig.D(i, i) = eig.D(i, i) + 3.0*frac*d;
  }

  M = eig.recompose();

  this->FromMatrix(M);
}
