//
// 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 <EdgeDetectAngle.hpp>

#include <cuda/CudaExport.hpp>
#include <ITKConfig.hpp>
#include <ITKFilter.hpp>

#include <iostream>

#include <itkCenteredAffineTransform.h>
#include <itkImportImageFilter.h>
#include <itkLinearInterpolateImageFunction.h>
#include <itkMinimumImageFilter.h>
#include <itkResampleImageFilter.h>

#include <StackProcessMorphology.hpp>
#include <MeshProcessHeatMap.hpp>

using namespace std;

namespace mgx
{
  mgx_EXPORT void getPointList(const Mesh* mesh, HVec3F& pts, HVec3F& nrmls);
  mgx_EXPORT void getBBox(const HVec3F& pts, Point3f* bBox);
  
  // Get corners of bounding box
  void getBBoxCorners(Point3f* bBox, HVec3F& bBoxCorners)
  {
    bBoxCorners.clear();
    for(int i = 0; i < 4; i++)
      bBoxCorners.push_back(bBox[0]);
  
    bBoxCorners[1].x() = bBox[1].x();
    bBoxCorners[2].y() = bBox[1].y();
    bBoxCorners[3].z() = bBox[1].z();
  
    for(int i = 0; i < 4; i++)
      bBoxCorners.push_back(bBox[1]);
  
    bBoxCorners[5].x() = bBox[0].x();
    bBoxCorners[6].y() = bBox[0].y();
    bBoxCorners[7].z() = bBox[0].z();
  }
  
  bool EdgeDetectAngle::run(Stack* s1, Store*, Store* output, float threshold, float multiplier,
                                   float factor, uint fillValue, bool fillHoles)
  {
    // current stack
    // main store = original store
    // work store = contains existing edge detect result
  
    typedef itk::ResampleImageFilter<UImageType, UImageType> FilterType;
    typedef itk::CenteredAffineTransform<double, 3> TransformType;
    typedef itk::LinearInterpolateImageFunction<UImageType, double> InterpolatorType;
    typedef itk::MinimumImageFilter<UImageType> MinimumImageFilter;
    typedef itk::ImportImageFilter<unsigned short int, 3> ImportFilterType;
    typedef UImageType::PixelType PixelType;
  
    // backup work store (to keep the old edge detect result)
    HVecUS backupData(s1->work()->data());
    Point3f backupStep(s1->step());
    Point3f backupOrigin(s1->origin());
    Point3f backupSize(s1->size());
  
    // get libqgl position and transformation matrices of the frames
    Matrix4d mGL, mGLinv;
    s1->getFrame().getMatrix(mGL.data());
    mGLinv = inverse(mGL);
  
    // Get bounding box
    HVec3F pts1;
    Point3f bBox[2];
    HVec3F bBoxCorners, bBoxCornersRot;
  
    // obtain bbox
    pts1.clear();
    pts1.push_back(s1->imageToWorld(Point3i(0, 0, 0)));
    pts1.push_back(s1->imageToWorld(s1->size()));
    getBBox(pts1, bBox);
  
    // Get all corners of bbox
    getBBoxCorners(bBox, bBoxCorners);
  
    // find xLow, xHigh, yLow, yHigh, zLow, zHigh
    Point3f low, high;
  
    for(int i = 0; i < 3; i++) {
      low[i] = 1E20;
      high[i] = -1E20;
    }
  
    // rotate bbox
    for(int j = 0; j < 8; j++) {
      Point3f currentCorner(multMatrix4Point3(mGLinv, Point3d(bBoxCorners[j])));
      for(int i = 0; i < 3; i++) {
        if(currentCorner[i] < low[i])
          low[i] = currentCorner[i];
        if(currentCorner[i] > high[i])
          high[i] = currentCorner[i];
      }
    }
  
    // increase the size
    Point3u sz = s1->size();
    Point3f voxelSize = s1->step();
    Point3f sizeDifNeg;   // for originMove
  
    for(int i = 0; i < 3; i++) {
      sizeDifNeg[i] = 0;
    }
  
    for(int i = 0; i < 3; i++)
      sizeDifNeg[i] += (-low[i] + bBox[0][i]) * voxelSize[i];
  
    itk::Size<3> itkSize;
  
    for(int i = 0; i < 3; i++) {
      itkSize[i] = ceil((high[i] - low[i]) / voxelSize[i]);
    }
    // correct orgin (due to increased size works in only one direction)
    Point3f origin = s1->origin();
    Point3f originN;
  
    for(int i = 0; i < 3; i++)
      originN[i] = origin[i] - sizeDifNeg[i] / voxelSize[i];
  
    // rotation and translation into ITK
    TransformType::Pointer transform = TransformType::New();
    TransformType::ParametersType parameters = transform->GetParameters();
  
    // put rotation matrix in parameters
    for(int i = 0; i < 3; i++)
      for(int j = 0; j < 3; j++)
        parameters[i * 3 + j] = mGL[i][j];
  
    // put corrected rotation center in parameters (fix of the "1-pixel-problem")
    double correction = -0.5f;
    for(int i = 0; i < 3; i++)
      parameters[3 * 3 + i] = correction * voxelSize[i];
  
    // put translation vector in parameters
    for(int i = 0; i < 3; i++)
      parameters[4 * 3 + i] = mGL[i][3];
  
    transform->SetParameters(parameters);
    UImageConverter::Pointer converter = UImageConverter::New();
    converter->SetStore(s1->main());
    // filter
    FilterType::Pointer filter = FilterType::New();
    filter->SetInput(converter->GetOutput());
    filter->SetTransform(transform);
    filter->SetSize(itkSize);
    filter->SetOutputSpacing(Point3d(voxelSize).data());
  
    filter->SetOutputOrigin(Point3d(originN).data());
  
    InterpolatorType::Pointer interpolator = InterpolatorType::New();
    filter->SetInterpolator(interpolator);
    filter->SetDefaultPixelValue(0);
    filter->Update();
  
    // save result in temp store
    HVecUS tempVec(size_t(itkSize[0] * itkSize[1] * itkSize[2]));
  
    // copy filter output to temp vector
    UImageType::RegionType region = filter->GetOutput()->GetBufferedRegion();
  
    ushort* out = tempVec.data();
  
    typedef itk::ImageRegionConstIterator<UImageType> IteratorType;
  
    IteratorType it(filter->GetOutput(), region);
  
    it.GoToBegin();
    while(!it.IsAtEnd()) {
      *out++ = it.Get();
      ++it;
    }
  
    // do edge detect from top (identical to EdgeDetect)
    ushort lowthresh = trim(ushort(float(threshold) * factor), ushort(0), ushort(0xFFFF));
    ushort fill = trim(fillValue, uint(0), uint(0xFFFF));
    Point3i size(itkSize[0], itkSize[1], itkSize[2]);   // s1->size();
    edgeDetectGPU(size, lowthresh, threshold, multiplier, fill, tempVec, tempVec);

    // Do fill holes if required
    if(fillHoles) {
		
		
		
     FillHolesProcess *fillHolesProcess;
     if(!getProcess("Stack/Morphology/Fill Holes", fillHolesProcess))  
      throw(QString("EdgeDetectAngle: Unable to make FillHolesProcess process"));       		

     fillHolesProcess->run(tempVec, tempVec, size, fillHolesProcess->parm("X Radius").toUInt(), 
                          fillHolesProcess->parm("Y Radius").toUInt(), 1, fillHolesProcess->parm("Fill Value").toUInt(), fillValue);
    }
  
    // rotate work stack back to original position
    TransformType::Pointer transformB = TransformType::New();
    TransformType::ParametersType parametersB = transform->GetParameters();
  
    // put rotation matrix in parameters
    for(int i = 0; i < 3; i++)
      for(int j = 0; j < 3; j++) {
        parametersB[i * 3 + j] = mGLinv[i][j];
      }
  
    // put corrected rotation center in parameters
    for(int i = 0; i < 3; i++)
      parametersB[3 * 3 + i] = correction * voxelSize[i];
  
    // put translation vector in parameters
    for(int i = 0; i < 3; i++)
      parametersB[4 * 3 + i] = mGLinv[i][3];
  
    transformB->SetParameters(parametersB);
  
    // put new edge detection image + size/spacing/origin info into importer
    ImportFilterType::Pointer importer = ImportFilterType::New();
  
    UImageType::SizeType size2;
    size2[0] = itkSize[0];
    size2[1] = itkSize[1];
    size2[2] = itkSize[2];
  
    UImageType::IndexType start;
    start[0] = start[1] = start[2] = 0;
  
    UImageType::RegionType region2;
    region2.SetSize(size2);
    region2.SetIndex(start);
  
    importer->SetRegion(region2);
  
    double spacing[3];
    spacing[0] = voxelSize[0];
    spacing[1] = voxelSize[1];
    spacing[2] = voxelSize[2];
  
    importer->SetSpacing(spacing);
  
    double origin2[3];
    origin2[0] = originN[0];
    origin2[1] = originN[1];
    origin2[2] = originN[2];
  
    importer->SetOrigin(origin2);
  
    const bool importFilterWillDeleteTheInputBuffer = false;
    PixelType* pixelData = const_cast<PixelType*>(tempVec.data());
    unsigned long totalNumberOfPixels = itkSize[0] * itkSize[1] * itkSize[2];
  
    importer->SetImportPointer(pixelData, totalNumberOfPixels, importFilterWillDeleteTheInputBuffer);
    importer->Update();
  
    // input for first new filter
    FilterType::Pointer filterB = FilterType::New();
    filterB->SetInput(importer->GetOutput());
    filterB->SetTransform(transformB);
  
    // original size
    itk::Size<3> itkSizeB;
    for(int i = 0; i < 3; i++)
      itkSizeB[i] = sz[i];
  
    filterB->SetSize(itkSizeB);
    filterB->SetOutputSpacing(Point3d(voxelSize).data());
    filterB->SetOutputOrigin(Point3d(origin).data());
    filterB->SetInterpolator(interpolator);
    filterB->SetDefaultPixelValue(0);
    filterB->Update();
  
    // load data of the old edge detect result
    s1->setStep(backupStep);
    s1->setOrigin(backupOrigin);
    s1->setSize(Point3u(backupSize));
    s1->work()->data() = backupData;
  
    // and put its data into a second new filter
    UImageConverter::Pointer converterB2 = UImageConverter::New();
    converterB2->SetStore(s1->work());
    FilterType::Pointer filterB2 = FilterType::New();
    filterB2->SetInput(converterB2->GetOutput());
    filterB2->SetSize(itkSizeB);
    filterB2->SetOutputSpacing(Point3d(voxelSize).data());
    filterB2->SetOutputOrigin(Point3d(origin).data());
    filterB2->Update();
    filterB2->SetInterpolator(interpolator);
    filterB2->SetDefaultPixelValue(0);
  
    // combine the two edge detect filters to get the result
    MinimumImageFilter::Pointer minImageFilter = MinimumImageFilter::New();
    minImageFilter->SetInput(0, filterB2->GetOutput());
    minImageFilter->SetInput(1, filterB->GetOutput());
    minImageFilter->Update();
  
    // Stack *s2 = stack(1);
    UImageConverter::TransferImage(s1, output, minImageFilter->GetOutput());
    // UImageConverter::TransferImage(s2, s2->work(), filter->GetOutput());
    output->changed();
  
    return true;
  }
  
  REGISTER_PROCESS(EdgeDetectAngle);
}
