
#include "DTImageIsospectralFilter.h"

#include <cmath>
#include <iostream>

#include <float.h>

DTImageIsospectralFilter::
DTImageIsospectralFilter()
{
  m_TimeStep = 1e-3;
  m_NumberOfIterations = 100;
}

DTImageIsospectralFilter::
~DTImageIsospectralFilter()
{

}

void
DTImageIsospectralFilter::
Update()
{

  m_Output = 0;

  if (m_Input.IsNull())
    return;

  // Copy input data to output data
  DTImageRegionType region = m_Input->GetLargestPossibleRegion();

  DTImageSizeType size = region.GetSize();

  m_Output = DTImageType::New();
  m_Output->CopyInformation(m_Input);
  m_Output->SetRegions(region);
  m_Output->Allocate();

  // Initialize output with input values
  DTImageIndexType ind;
  for (ind[2] = 0; ind[2] < (long)size[2]; ind[2]++)
    for (ind[1] = 0; ind[1] < (long)size[1]; ind[1]++)
      for (ind[0] = 0; ind[0] < (long)size[0]; ind[0]++)
      {
        m_Output->SetPixel(ind, m_Input->GetPixel(ind));
      }

  // Allocate conductivity image
  m_Internal = InternalImageType::New();
  m_Internal->CopyInformation(m_Input);
  m_Internal->SetRegions(region);
  m_Internal->Allocate();
  m_Internal->FillBuffer(0);

  // Flow one time step
  float t = 0;
  for (unsigned int k = 0; k < m_NumberOfIterations; k++)
  {
    itkDebugMacro(<< "Step, t = " << t);
    this->Step();
    t += m_TimeStep;
  }

  m_Internal = 0;

}

void DTImageIsospectralFilter::
Step()
{

  DTImageRegionType region = m_Input->GetLargestPossibleRegion();

  DTImageSizeType size = region.GetSize();
  DTImageSpacingType spacing = m_Input->GetSpacing();

  DTImageIndexType ind;
  DTImageIndexType f_ind;
  DTImageIndexType b_ind;

  itkDebugMacro(<< "Compute conductivity");
  // Internal image at x is phi_prime(s)/s
  // where s is grad mag
  for (ind[2] = 1; ind[2] < (size[2]-1); ind[2]++)
    for (ind[1] = 1; ind[1] < (size[1]-1); ind[1]++)
      for (ind[0] = 1; ind[0] < (size[0]-1); ind[0]++)
      {
        // Gradient
        DiffusionTensor dX[3];

        f_ind[0] = ind[0]+1;
        f_ind[1] = ind[1];
        f_ind[2] = ind[2];
        b_ind[0] = ind[0]-1;
        b_ind[1] = ind[1];
        b_ind[2] = ind[2];
        dX[0] = m_Output->GetPixel(f_ind) - m_Output->GetPixel(b_ind);
        dX[0] /= spacing[0];

        f_ind[0] = ind[0];
        f_ind[1] = ind[1]+1;
        f_ind[2] = ind[2];
        b_ind[0] = ind[0];
        b_ind[1] = ind[1]-1;
        b_ind[2] = ind[2];
        dX[1] = m_Output->GetPixel(f_ind) - m_Output->GetPixel(b_ind);
        dX[1] /= spacing[1];

        f_ind[0] = ind[0];
        f_ind[1] = ind[1];
        f_ind[2] = ind[2]+1;
        b_ind[0] = ind[0];
        b_ind[1] = ind[1];
        b_ind[2] = ind[2]-1;
        dX[2] = m_Output->GetPixel(f_ind) - m_Output->GetPixel(b_ind);
        dX[2] /= spacing[2];

        // Compute grad mag, sqrt of sum of trace of d_iX'*d_iX
        // |dX| = sqrt sum_i{trace(d_iX'*d_iX)}
        double s = FLT_EPSILON;
        for (unsigned int i = 0; i < 3; i++)
        {
          MatrixType A = dX[i].GetMatrix();

          MatrixType AtA = A.transpose() * A;

          double trace = 0;
          for (unsigned int k = 0; k < 3; k++)
            trace += AtA(k, k);

          s += trace;
        }
        s = sqrt(s);

        // phi(s) = 2*sqrt(1+s^2) - 2
        // phi_prime(s) = 4*s / sqrt(1+s^2)

        double phi_prime = 4*s / sqrt(1 + s*s);

        m_Internal->SetPixel(ind, phi_prime / s);
      }

  // Flow
  for (ind[2] = 1; ind[2] < (size[2]-1); ind[2]++)
    for (ind[1] = 1; ind[1] < (size[1]-1); ind[1]++)
      for (ind[0] = 1; ind[0] < (size[0]-1); ind[0]++)
      {
        DiffusionTensor center = m_Output->GetPixel(ind);

        f_ind[0] = ind[0]+1;
        f_ind[1] = ind[1];
        f_ind[2] = ind[2];
        b_ind[0] = ind[0]-1;
        b_ind[1] = ind[1];
        b_ind[2] = ind[2];
        DiffusionTensor flowx =
          (m_Output->GetPixel(f_ind) - center) * m_Internal->GetPixel(f_ind)
          - 
          (center - m_Output->GetPixel(b_ind)) * m_Internal->GetPixel(b_ind);
        flowx /= spacing[0];

        f_ind[0] = ind[0];
        f_ind[1] = ind[1]+1;
        f_ind[2] = ind[2];
        b_ind[0] = ind[0];
        b_ind[1] = ind[1]-1;
        b_ind[2] = ind[2];
        DiffusionTensor flowy =
          (m_Output->GetPixel(f_ind) - center) * m_Internal->GetPixel(f_ind)
          - 
          (center - m_Output->GetPixel(b_ind)) * m_Internal->GetPixel(b_ind);
        flowy /= spacing[1];

        f_ind[0] = ind[0];
        f_ind[1] = ind[1];
        f_ind[2] = ind[2]+1;
        b_ind[0] = ind[0];
        b_ind[1] = ind[1];
        b_ind[2] = ind[2]-1;
        DiffusionTensor flowz =
          (m_Output->GetPixel(f_ind) - center) * m_Internal->GetPixel(f_ind)
          - 
          (center - m_Output->GetPixel(b_ind)) * m_Internal->GetPixel(b_ind);
        flowz /= spacing[2];

        //G(X)(p) = -div{phi_prime(p) * d(X(p))}
        DiffusionTensor::MatrixType G =
          flowx.GetMatrix() + flowy.GetMatrix() + flowz.GetMatrix();
        G *= -1.0;

        DiffusionTensor::MatrixType GpGt = G + G.transpose();
        GpGt *= -1.0;

        // Note: [A, B] = AB - BA
        // Isospectral flow follows
        // dI/dt = [I, [I, -(G+G')]]
        DiffusionTensor::MatrixType I = center.GetMatrix();
        DiffusionTensor::MatrixType P = I*GpGt - GpGt*I;
        DiffusionTensor::MatrixType dI = I*P - P*I;

        MatrixType newI = I + dI*m_TimeStep;

        DiffusionTensor T;
        T.FromMatrix(newI);
        m_Output->SetPixel(ind, T);
      }

}
