//
// This file is part of MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2016 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 <StackProcessFilters.hpp>
#include <SystemProcessLoad.hpp>

#include <CImg.h>

using namespace cimg_library;

typedef CImg<ushort> CImgUS;

namespace mgx 
{
  
  REGISTER_PROCESS(StackSwapBytes);
  
  bool AverageStack::run(const Store* input, Store* output, Point3i radius, uint steps)
  {
    if(steps <= 0)
      throw std::string("Average voxels, steps must be > 0");
    const Stack* stack = output->stack();
    SETSTATUS("Average( " << radius << ", " << steps << ")");
    progressStart(QString("Average pixels for Stack %1").arg(stack->userId()), 100);
    const HVecUS& data = input->data();
    HVecUS& outputData = output->data();
    Point3i imgSize(stack->size());
    if(averageGPU(imgSize, radius, data, outputData)) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    for(uint i = 1; i < steps; i++) {
      if(!progressAdvance(1))
        userCancel();
      if(averageGPU(imgSize, radius, outputData, outputData)) {
        setErrorMessage("Error while running CUDA process");
        return false;
      }
    }
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(AverageStack);
  
  bool MedianStack::run(const Store* input, Store* output, Point3i radius)
  {
    const Stack* stack = output->stack();
    SETSTATUS("Median( " << radius << ")");
    progressStart(QString("Median pixels for Stack %1").arg(stack->userId()), 100);
    const HVecUS& data = input->data();
    HVecUS& outputData = output->data();
    Point3i imgSize(stack->size());
    if(medianGPU(imgSize, radius, data, outputData)) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(MedianStack);
  
  bool MedianOfMediansStack::run(const Store* input, Store* output, Point3i radius)
  {
    const Stack* stack = output->stack();
    SETSTATUS("Median-of-medians( " << radius << ")");
    progressStart(QString("Median-of-median pixels for Stack %1").arg(stack->userId()), 100);
    const HVecUS& data = input->data();
    HVecUS& outputData = output->data();
    Point3i imgSize(stack->size());
    if(medianOfMediansGPU(imgSize, radius, data, outputData)) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(MedianOfMediansStack);
  
  bool BinarizeStack::run(const Store* input, Store* output, ushort threshold)
  {
    const HVecUS& idata = input->data();
    HVecUS& data = output->data();
    progressStart("Binarize stack", 100);
    #pragma omp parallel for
    for(ulong idx = 0; idx < data.size(); idx++) {
      data[idx] = (idata[idx] > threshold) ? 65535 : 0;
      if(idx % 100 == 0 and !progressAdvance(idx / 100))
        userCancel();
    }
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(BinarizeStack);
  
  bool BrightenStack::run(const Store* input, Store* output, float brightness)
  {
    if(input != output)
      output->data() = input->data();
    HVecUS& data = output->data();
    progressStart("Brighten/Darken stack", 100);
    #pragma omp parallel for
    for(ulong idx = 0; idx < data.size(); idx++) {
      float pix = float(data[idx]) * brightness;
      if(pix < 0)
        pix = 0;
      else if(pix > 0xFFFF)
        pix = 0xFFFF;
      data[idx] = ushort(pix);
      if(idx % 100 == 0 and !progressAdvance(idx / 100))
        userCancel();
    }
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(BrightenStack);
  
  bool ColorGradient::run(const Store* input, Store* output, float colorGradDiv)
  {
    Point3i imgSize(input->stack()->size());
    const HVecUS& data = input->data();
    HVecUS& workData = output->data();
    if(colorGradGPU(imgSize, colorGradDiv, data, workData)) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(ColorGradient);
  
  bool FilterStack::run(const Store* input, Store* output, uint lowFilter, uint highFilter)
  {
    const Stack* stack = output->stack();
    const HVecUS& data = input->data();
    HVecUS& workData = output->data();
    int highCnt = 0, lowCnt = 0;
  
    lowFilter = trim(lowFilter, uint(0), uint(0xFFFF));
    highFilter = trim(highFilter, uint(0), uint(0xFFFF));
  
    for(uint idx = 0; idx < data.size(); idx++) {
      if(highFilter > 0 and data[idx] > highFilter) {
        highCnt++;
        workData[idx] = highFilter;
      } else if(data[idx] < lowFilter) {
        lowCnt++;
        workData[idx] = 0;
      } else
        workData[idx] = data[idx];
    }
    SETSTATUS("Stack " << stack->userId() << " Filtered out " << lowCnt << " low pixels, and " << highCnt
                       << " high pixels");
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(FilterStack);
  
  bool InvertStack::run(const Store* input, Store* output)
  {
    const HVecUS& data = input->data();
    HVecUS& workData = output->data();
    for(uint i = 0; i < data.size(); i++)
      workData[i] = 0xFFFF - data[i];
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(InvertStack);
  
  bool GaussianBlurStack::run(const Store* input, Store* output, Point3f sigma)
  {
    const Stack* stack = input->stack();
    Point3i imgSize(stack->size());
    gaussianBlurGPU(imgSize, stack->step(), sigma, input->data(), output->data());
    output->copyMetaData(input);
    output->changed();
    return true;
  }
  REGISTER_PROCESS(GaussianBlurStack);

  bool DiffGaussians::run(const Store* input, Store* output, Point3f sigma1, Point3f sigma2)
  {
    const Stack* stack = input->stack();
    Point3i imgSize(stack->size());
    HVecUS output1(input->data());
    if(norm(sigma1) > 0)
      gaussianBlurGPU(imgSize, stack->step(), sigma1, input->data(), output1);

    HVecUS output2(input->data());
    if(norm(sigma2) > 0)
      gaussianBlurGPU(imgSize, stack->step(), sigma2, input->data(), output2);

    // Subtract Guassians 
    #pragma omp parallel for
    for(size_t i = 0; i < input->data().size(); ++i)
      output->data()[i] = trim(int(output1[i]) - int(output2[i]), 0, 65535);

    output->copyMetaData(input);
    output->changed();
    return true;
  }
  REGISTER_PROCESS(DiffGaussians); 

  bool SharpenStack::run(const Store* input, Store* output, 
                                       const Point3f &sigma, float amount, int passes)
  {
    const Stack* stack = input->stack();
    Point3i imgSize(stack->size());

    HVecUS src(input->data()), dst(input->data().size());
    HVecUS *srcP = &src, *dstP = &dst;
    progressStart("Sharpening Stack", passes);
    for(int i = 0; i < passes; ++i) {
      sharpenGPU(imgSize, stack->step(), sigma, amount, *srcP, *dstP);
      std::swap(src, dst);
      if(!progressAdvance(i + 1))
        userCancel();
    }
    output->copyMetaData(input);
    output->changed();
    output->data() = *srcP;

    return true;
  }
  REGISTER_PROCESS(SharpenStack);
  
  //bool ApplyKernelStack::run(const Store* input, Store* output, const HVecF& kernelX, const HVecF& kernelY,
                                    //const HVecF& kernelZ)
  //{
    //const Stack* stack = input->stack();
    //Point3i imgSize(stack->size());
  
    //if((kernelX.size() > 0 and kernelX.size() < 3)or (kernelY.size() > 0 and kernelY.size() < 3)
       //or (kernelZ.size() > 0 and kernelZ.size() < 3)) {
      //setErrorMessage("Kernel must be empty or at least 3 numbers long");
      //return false;
    //}
    //if((kernelX.size() > 0 and (kernelX.size() % 2) == 0)or (kernelY.size() > 0 and (kernelY.size() % 2) == 0)
       //or (kernelZ.size() > 0 and (kernelZ.size() % 2) == 0)) {
      //setErrorMessage("Kernel size must be odd");
      //return false;
    //}
    //applyKernelGPU(imgSize, kernelX, kernelY, kernelZ, input->data(), output->data());
    //output->copyMetaData(input);
    //output->changed();
    //return true;
  //}
  //REGISTER_PROCESS(ApplyKernelStack);
  
  bool NormalizeStack::run(const Store* input, Store* output, Point3f radius, Point3f sigma, uint thresh,
                                  float blurFactor)
  {
    if(blurFactor < 0 or blurFactor > 1)
      throw(QString("Blur Factor must be between 0 and 1"));
    if(thresh > 65535)
      throw(QString("Threshold must be between 0 and 65535"));
    if(radius.x() < 0 or radius.y() < 0 or radius.z() < 0)
      throw(QString("Radii for blurring must all be >= 0"));
  
    const Stack* stack = input->stack();
    Point3i imgSize(stack->size());
    const HVecUS& data = input->data();
  
    Point3i voxRadius(toVoxelsCeil(radius, stack->step()));
  
    HVecUS bDilate(input->data().size());
    HVecUS tDilate(input->data().size());
    HVecUS& workData = output->data();
  
    dilateGPU(imgSize, voxRadius, false, input->data(), tDilate);
    if(norm(sigma) == 0) {
      bDilate = tDilate;
      blurFactor = 0;
    } else {
      gaussianBlurGPU(imgSize, stack->step(), sigma, input->data(), bDilate);
      dilateGPU(imgSize, voxRadius, false, bDilate, bDilate);
    }
    if(blurFactor > 1.0)
      blurFactor = 1.0;
  
    // Blend the value based on the range and the blurred range
    for(uint idx = 0; idx < workData.size(); idx++)
      if(tDilate[idx] > thresh)
        workData[idx] = trim(int((blurFactor * trim(float(data[idx]) / bDilate[idx], 0.0f, 1.0f)
                                  + (1.0 - blurFactor) * data[idx] / tDilate[idx]) * 0xFFFF),
                             0, 0xFFFF);
      else
        workData[idx] = 0;
  
    output->copyMetaData(input);
    output->changed();
    return true;
  }
  
  REGISTER_PROCESS(NormalizeStack);
  
  // CImg filters
  bool CImgGaussianBlurStack::run(const Store* input, Store* output, uint radius)
  {
    if(input != output)
      output->data() = input->data();
    Point3u size = input->stack()->size();
    CImgUS image(output->data().data(), size.x(), size.y(), size.z(), 1, true);
    image.blur(radius);
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(CImgGaussianBlurStack);

  bool CImgMedianBlurStack::run(const Store* input, Store* output, uint radius)
  {
    if(input != output)
      output->data() = input->data();
    Point3u size = input->stack()->size();
    CImgUS image(output->data().data(), size.x(), size.y(), size.z(), 1, true);
    image.blur_median(radius);
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(CImgMedianBlurStack); 

  bool CImgLaplaceStack::run(const Store* input, Store* output)
  {
    if(input != output)
      memcpy(output->data().data(), input->data().data(), input->data().size() * 2);
    Point3u size = input->stack()->size();
    CImgUS image(output->data().data(), size.x(), size.y(), size.z(), 1, true);
    image.laplacian();
    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(CImgLaplaceStack);

  bool TopHatStack::run(Stack* stack, const Store* input, Store* output, 
            uint xradius, uint yradius, uint zradius, bool roundNhbd)
  {
    Point3i radius(xradius, yradius, zradius);
    Point3i size(stack->size());
    // Backup input image
    HVecUS origInput = input->data();
    // Erode and dilate (opening)
    if(erodeGPU(size, radius, input->labels(), roundNhbd, input->data(), output->data())) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    if(dilateGPU(size, radius, roundNhbd, output->data(), output->data())) {
      setErrorMessage("Error while running CUDA process");
      return false;
    }
    // Subtract opening from original image
    #pragma omp parallel for
    for(size_t i = 0; i < output->data().size(); ++i)
      output->data()[i] = trim(int(origInput[i]) - int(output->data()[i]), 0, 65535);

    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(TopHatStack); 

  bool AddNoiseStack::run(Stack* stack, const Store* input, Store* output, uint amount)
  {
    // Backup input image
    auto &inData = input->data();
    auto &outData = output->data();

    // Subtract opening from original image
    #pragma omp parallel for
    for(size_t i = 0; i < output->data().size(); ++i)
      outData[i] = trim(int(inData[i]) + int(double(rand())/double(RAND_MAX) * amount), 0, 65535);

    output->changed();
    output->copyMetaData(input);
    return true;
  }
  REGISTER_PROCESS(AddNoiseStack); 
}
