//
// This file is part of MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2015 Richard S. Smith and collaborators.
//
// If you use MorphoGraphX in your work, please cite:
//   http://dx.doi.org/10.7554/eLife.05864
//
// MorphoGraphX is free software, and is licensed under under the terms of the 
// GNU General (GPL) Public License version 2.0, http://www.gnu.org/licenses.
// 
#include "ITKRegistration.hpp"

#include "ITKProgress.hpp"

#include <itkDemonsRegistrationFilter.h>
#include <itkLevelSetMotionRegistrationFilter.h>
#include <itkHistogramMatchingImageFilter.h>
#include <itkCastImageFilter.h>
#include <itkWarpImageFilter.h>
#include <itkLinearInterpolateImageFunction.h>
#include <itkImageFileWriter.h>
#include <itkVTKImageIO.h>
#include <itkImportImageFilter.h>

#include <Information.hpp>

namespace mgx
{
  
  bool ITKDemonsRegistration::run(const Store* reference, Stack* image_stk, 
              QString output, bool thresholdAtMean, int nbHistLevels, int nbMatchPoints, 
              int nbIterations, float sigma)
  {
    Information::out << "Starting process. Output will be saved in '" << output << "'" << endl;
    Information::out << "Size of reference = " << reference->stack()->size() << endl;
    Information::out << "Size of image = " << reference->stack()->size() << endl;
  
    Store* image = image_stk->currentStore();
  
    const HVecUS& rdata = reference->data();
    const HVecUS& idata = image->data();
    ushort rmin = 0xFFFF, rmax = 0, imin = 0xFFFF, imax = 0;
  #pragma omp parallel for reduction(min : rmin) reduction(max : rmax)
    for(size_t i = 0; i < rdata.size(); ++i) {
      if(rdata[i] < rmin)
        rmin = rdata[i];
      if(rdata[i] > rmax)
        rmax = rdata[i];
    }
  
  #pragma omp parallel for reduction(min : imin) reduction(max : imax)
    for(size_t i = 0; i < idata.size(); ++i) {
      if(idata[i] < imin)
        imin = idata[i];
      if(idata[i] > imax)
        imax = idata[i];
    }
  
    Information::out << "Range data reference: " << rmin << "-" << rmax << endl;
    Information::out << "Range data image    : " << imin << "-" << imax << endl;
  
    UImageConverter::Pointer ref_converter = UImageConverter::New();
    ref_converter->SetStore(reference);
    UImageConverter::Pointer img_converter = UImageConverter::New();
    img_converter->SetStore(image);
  
    ref_converter->Update();
    img_converter->Update();
  
    typedef itk::Image<float, 3> InternalImageType;
  
    typedef itk::CastImageFilter<UImageType, InternalImageType> ImageCasterType;
  
    ImageCasterType::Pointer referenceCaster = ImageCasterType::New();
    ImageCasterType::Pointer imgCaster = ImageCasterType::New();
  
    referenceCaster->SetInput(ref_converter->GetOutput());
    imgCaster->SetInput(img_converter->GetOutput());
  
    typedef itk::HistogramMatchingImageFilter<InternalImageType, InternalImageType> MatchingFilterType;
  
    MatchingFilterType::Pointer matcher = MatchingFilterType::New();
    matcher->SetInput(imgCaster->GetOutput());
    matcher->SetReferenceImage(referenceCaster->GetOutput());
  
    matcher->SetNumberOfHistogramLevels(nbHistLevels);
    matcher->SetNumberOfMatchPoints(nbMatchPoints);
    matcher->SetThresholdAtMeanIntensity(thresholdAtMean);
  
    typedef itk::Vector<float, 3> VectorPixelType;
    typedef itk::Image<VectorPixelType, 3> DeformationFieldType;
  
    // Compute initial displacement field
    typedef itk::ImportImageFilter<VectorPixelType, 3> ImportDeformationFieldType;
    ImportDeformationFieldType::Pointer initDeformation = ImportDeformationFieldType::New();
  
    initDeformation->SetSpacing(img_converter->GetSpacing());
    initDeformation->SetRegion(img_converter->GetRegion());
    initDeformation->SetOrigin(img_converter->GetOrigin());
  
    // No initial deformation, as anyway it seems to make the algorithm crash
    /*
     *  std::vector<VectorPixelType> initDefData(image->size());
     *
     *  Matrix4f trans;
     *  {
     *    qglviewer::Frame f1 = image_stk->getFrame().inverse();
     *    qglviewer::Frame f2 = reference->stack()->getFrame();
     *    f2.setReferenceFrame(&f1);
     *    double m[16];
     *    f2.getWorldMatrix(m);
     *    trans = Matrix4f(m, GL_STYLE);
     *  }
     *
     *  {
     *    Point3f start = image_stk->origin();
     *    Point3f step = image_stk->step();
     *    Point3u size = image_stk->size();
     **#pragma omp parallel for
     *    for(uint z = 0 ; z < size.z() ; ++z)
     *    {
     *      uint k = image_stk->offset(0u, 0u, z);
     *      for(uint y = 0 ; y < size.y() ; ++y)
     *        for(uint x = 0 ; x < size.x() ; ++x, ++k)
     *        {
     *          Point4f pt(image_stk->imageToWorld(Point3i(x,y,z)));
     *          pt.t() = 1.;
     *          Point4f tr = trans * pt;
     *          tr /= tr.t();
     *          Point3f diff = Point3f(tr - pt);
     *          VectorPixelType px;
     *          px[0] = diff.x();
     *          px[1] = diff.y();
     *          px[2] = diff.z();
     *          initDefData[k] = px;
     *        }
     *    }
     *  }
     */
  
    ITKProgress progress("Demons Registration");
  
    // initDeformation->SetImportPointer(initDefData.data(), initDefData.size(), false);
    // initDeformation->Update();
  
    // And continue ...
    // typedef itk::DemonsRegistrationFilter<InternalImageType,InternalImageType,DeformationFieldType>
    // RegistrationFilterType;
    typedef itk::LevelSetMotionRegistrationFilter<InternalImageType, InternalImageType, DeformationFieldType>
      RegistrationFilterType;
  
    RegistrationFilterType::Pointer filter = RegistrationFilterType::New();
    filter->SetFixedImage(referenceCaster->GetOutput());
    filter->SetMovingImage(matcher->GetOutput());
  
    filter->SetNumberOfIterations(nbIterations);
    filter->SetStandardDeviations(sigma);
  
    // filter->SetInitialDisplacementField(initDeformation->GetOutput());
  
    progress.setFilter(filter);
  
    filter->Update();
  
    typedef itk::WarpImageFilter<UImageType, UImageType, DeformationFieldType> WarperType;
    typedef itk::LinearInterpolateImageFunction<UImageType, double> InterpolationType;
  
    UImageType::Pointer ref = ref_converter->GetOutput();
    DeformationFieldType::Pointer deform = filter->GetOutput();
    const itk::Vector<float, 3>* ddata = deform->GetBufferPointer();
  
    DeformationFieldType::RegionType dregion = deform->GetBufferedRegion();
    DeformationFieldType::RegionType::SizeType defSize = dregion.GetSize();
  
    long unsigned int dsize = defSize[0] * defSize[1] * defSize[2];
    float xmin = 1e10, ymin = 1e10, zmin = 1e10;
    float xmax = -1e10, ymax = -1e10, zmax = -1e10;
  #pragma omp parallel for reduction(min : xmin, ymin, zmin) reduction(max : xmax, ymax, zmax)
    for(long unsigned int i = 0; i < dsize; ++i) {
      const itk::Vector<float, 3>& value = ddata[i];
      if(value[0] < xmin)
        xmin = value[0];
      if(value[0] > xmax)
        xmax = value[0];
      if(value[1] < ymin)
        ymin = value[1];
      if(value[1] > ymax)
        ymax = value[1];
      if(value[2] < zmin)
        zmin = value[2];
      if(value[2] > zmax)
        zmax = value[2];
    }
    Information::out << "Range of deformations: (" << xmin << "," << ymin << "," << zmin << ")-(" << xmax << "," << ymax
                     << "," << zmax << ")" << endl;
    {
      typedef itk::ImageRegionConstIterator<DeformationFieldType> IteratorType;
  
      size_t cnt = 0;
      IteratorType it(deform, dregion);
  
      it.GoToBegin();
      while(!it.IsAtEnd()) {
        ++cnt;
        const itk::Vector<float, 3>& value = it.Get();
        if(value[0] < xmin)
          xmin = value[0];
        if(value[0] > xmax)
          xmax = value[0];
        if(value[1] < ymin)
          ymin = value[1];
        if(value[1] > ymax)
          ymax = value[1];
        if(value[2] < zmin)
          zmin = value[2];
        if(value[2] > zmax)
          zmax = value[2];
        ++it;
      }
  
      Information::out << "Visited " << cnt << " pixels" << endl;
    }
  
    Information::out << "Size of the deformation field: " << defSize[0] << "x" << defSize[1] << "x" << defSize[2]
                     << endl;
    Information::out << "Range of deformations: (" << xmin << "," << ymin << "," << zmin << ")-(" << xmax << "," << ymax
                     << "," << zmax << ")" << endl;
  
    WarperType::Pointer warper = WarperType::New();
    InterpolationType::Pointer interpolator = InterpolationType::New();
  
    warper->SetInput(img_converter->GetOutput());
    warper->SetInterpolator(interpolator);
    warper->SetOutputSpacing(ref->GetSpacing());
    warper->SetOutputOrigin(ref->GetOrigin());
    warper->SetDisplacementField(filter->GetOutput());
  
    Information::out << "Spacing = " << ref->GetSpacing()[0] << "," << ref->GetSpacing()[1] << ","
                     << ref->GetSpacing()[2] << endl;
  
    Information::out << "Origin = " << ref->GetOrigin()[0] << "," << ref->GetOrigin()[1] << "," << ref->GetOrigin()[2]
                     << endl;
  
    progress.setFilter(warper);
  
    warper->Update();
  
    if(not output.isEmpty()) {
      Information::out << "Saving file '" << output.toLocal8Bit() << "'" << endl;
      typedef itk::ImageFileWriter<DeformationFieldType> DeformationFieldWriter;
      DeformationFieldWriter::Pointer writer = DeformationFieldWriter::New();
  
      writer->SetInput(filter->GetOutput());
      writer->SetFileName(output.toLocal8Bit());
  
      itk::VTKImageIO::Pointer ImageIO = itk::VTKImageIO::New();
      writer->SetImageIO(ImageIO);
  
      ITKProgress progress("Saving ...");
      progress.setFilter(writer);
  
      writer->Write();
    }
  
    UImageType::ConstPointer result = warper->GetOutput();
    if(!UImageConverter::TransferImage(image_stk, image, result))
      return setErrorMessage("ITK Process tried to change the size of the stack");
    image->changed();
    return true;
  }
  
  REGISTER_PROCESS(ITKDemonsRegistration);
}
