/*============================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center (DKFZ)
All rights reserved.

Use of this source code is governed by a 3-clause BSD license that can be
found in the LICENSE file.

============================================================================*/

#include <mitkShowSegmentationAsSurface.h>
#include <mitkManualSegmentationToSurfaceFilter.h>
#include <mitkVtkRepresentationProperty.h>
#include <mitkCoreObjectFactory.h>
#include <mitkLabelSetImage.h>
#include <mitkLabelSetImageConverter.h>
#include <vtkPolyDataNormals.h>

#ifdef MITK_USE_OpenMP
#include <omp.h>
#endif

namespace mitk
{
  ShowSegmentationAsSurface::ShowSegmentationAsSurface()
    : m_UIDGeneratorSurfaces("Surface_"),
      m_IsLabelSetImage(false)
  {
  }

  ShowSegmentationAsSurface::~ShowSegmentationAsSurface()
  {
  }

  void ShowSegmentationAsSurface::Initialize(const NonBlockingAlgorithm *other)
  {
    Superclass::Initialize(other);

    bool syncVisibility(false);

    if (other)
    {
      other->GetParameter("Sync visibility", syncVisibility);
    }

    SetParameter("Sync visibility", syncVisibility);
    SetParameter("Median kernel size", 3u);
    SetParameter("Apply median", true);
    SetParameter("Smooth", true);
    SetParameter("Gaussian SD", 1.5);
    SetParameter("Decimate mesh", true);
    SetParameter("Decimation rate", 0.8);
    SetParameter("Wireframe", false);

    m_SurfaceNodes.clear();
  }

  bool ShowSegmentationAsSurface::ReadyToRun()
  {
    try
    {
      Image::Pointer image;
      GetPointerParameter("Input", image);
      MultiLabelSegmentation::Pointer labelSetImage;
      GetPointerParameter("Input", labelSetImage);

      return (labelSetImage.IsNotNull() || image.IsNotNull()) && GetGroupNode();
    }
    catch (std::invalid_argument &)
    {
      return false;
    }
  }

  bool ShowSegmentationAsSurface::ThreadedUpdateFunction()
  {
    bool smooth(true);
    GetParameter("Smooth", smooth);

    bool applyMedian(true);
    GetParameter("Apply median", applyMedian);

    bool decimateMesh(true);
    GetParameter("Decimate mesh", decimateMesh);

    unsigned int medianKernelSize(3);
    GetParameter("Median kernel size", medianKernelSize);

    double gaussianSD(1.5);
    GetParameter("Gaussian SD", gaussianSD);

    double reductionRate(0.8);
    GetParameter("Decimation rate", reductionRate);

    MITK_INFO << "Creating polygon model with smoothing " << smooth << " gaussianSD " << gaussianSD << " median "
              << applyMedian << " median kernel " << medianKernelSize << " mesh reduction " << decimateMesh
              << " reductionRate " << reductionRate;

    MultiLabelSegmentation::Pointer labelSetImage;
    GetPointerParameter("Input", labelSetImage);

    if (nullptr != labelSetImage)
    {
      const auto labels = labelSetImage->GetLabels();

      int numLabels = static_cast<int>(labels.size());
      m_SurfaceNodes.reserve(numLabels);

#ifdef MITK_USE_OpenMP
      omp_lock_t lock;
      omp_init_lock(&lock);
#endif

      #pragma omp parallel for
      for (int i = 0; i < numLabels; ++i)
      {
        auto labelImage = CreateLabelMask(labelSetImage, labels[i]->GetValue());

        if (labelImage.IsNull())
          continue;

        auto labelSurface = this->ConvertBinaryImageToSurface(labelImage);

        if (labelSurface.IsNull())
          continue;

        auto* polyData = labelSurface->GetVtkPolyData();

        if (smooth && (polyData->GetNumberOfPoints() < 1 || polyData->GetNumberOfCells() < 1))
        {
          MITK_WARN << "Label \"" << labels[i]->GetName() << "\" didn't produce any smoothed surface data (try again without smoothing).";
          continue;
        }

        auto node = DataNode::New();
        node->SetData(labelSurface);
        node->SetColor(labels[i]->GetColor());
        node->SetName(labels[i]->GetName());

#ifdef MITK_USE_OpenMP
        omp_set_lock(&lock);
#endif
        m_SurfaceNodes.push_back(node);
#ifdef MITK_USE_OpenMP
        omp_unset_lock(&lock);
#endif
      }

#ifdef MITK_USE_OpenMP
      omp_destroy_lock(&lock);
#endif
    }
    else
    {
      Image::Pointer image;
      GetPointerParameter("Input", image);
      auto surface = this->ConvertBinaryImageToSurface(image);

      if (surface.IsNotNull())
      {
        auto* polyData = surface->GetVtkPolyData();

        if (smooth && (polyData->GetNumberOfPoints() < 1 || polyData->GetNumberOfCells() < 1))
        {
          MITK_WARN << "Could not produce smoothed surface data (try again without smoothing).";
        }
        else
        {
          auto node = DataNode::New();
          node->SetData(surface);
          m_SurfaceNodes.push_back(node);
        }
      }
    }

    m_IsLabelSetImage = nullptr != labelSetImage;
    return true;
  }

  void ShowSegmentationAsSurface::ThreadedUpdateSuccessful()
  {
    for (const auto &node : m_SurfaceNodes)
    {
      bool wireframe = false;
      GetParameter("Wireframe", wireframe);

      if (wireframe)
      {
        auto representation = dynamic_cast<VtkRepresentationProperty *>(node->GetProperty("material.representation"));
        if (nullptr != representation)
          representation->SetRepresentationToWireframe();
      }

      node->SetProperty("opacity", FloatProperty::New(0.3f));
      node->SetProperty("line width", FloatProperty::New(1.0f));
      node->SetProperty("scalar visibility", BoolProperty::New(false));

      auto name = node->GetName();
      auto groupNode = this->GetGroupNode();

      if (!m_IsLabelSetImage)
      {
        if ((name.empty() || DataNode::NO_NAME_VALUE() == name) && nullptr != groupNode)
          name = groupNode->GetName();

        if (name.empty())
          name = "Surface";
      }

      bool smooth = true;
      GetParameter("Smooth", smooth);

      if (smooth)
        name.append(" (smoothed)");

      node->SetName(name);

      if (!m_IsLabelSetImage)
      {
        auto colorProp = groupNode->GetProperty("color");

        if (nullptr != colorProp)
        {
          node->ReplaceProperty("color", colorProp->Clone());
        }
        else
        {
          node->SetProperty("color", ColorProperty::New(1.0, 1.0, 0.0));
        }
      }

      bool showResult = true;
      GetParameter("Show result", showResult);

      bool syncVisibility = false;
      GetParameter("Sync visibility", syncVisibility);

      auto visibleProp = groupNode->GetProperty("visible");

      if (nullptr != visibleProp && syncVisibility)
      {
        node->ReplaceProperty("visible", visibleProp->Clone());
      }
      else
      {
        node->SetProperty("visible", BoolProperty::New(showResult));
      }

      if (!m_IsLabelSetImage)
      {
        Image::Pointer image;
        GetPointerParameter("Input", image);

        if (image.IsNotNull())
        {
          auto organTypeProp = image->GetProperty("organ type");

          if (nullptr != organTypeProp)
            node->GetData()->SetProperty("organ type", organTypeProp);
        }
      }

      this->InsertBelowGroupNode(node);
    }

    Superclass::ThreadedUpdateSuccessful();
  }

  Surface::Pointer ShowSegmentationAsSurface::ConvertBinaryImageToSurface(Image::Pointer binaryImage)
  {
    bool smooth = true;
    GetParameter("Smooth", smooth);

    bool applyMedian = true;
    GetParameter("Apply median", applyMedian);

    bool decimateMesh = true;
    GetParameter("Decimate mesh", decimateMesh);

    unsigned int medianKernelSize = 3;
    GetParameter("Median kernel size", medianKernelSize);

    double gaussianSD = 1.5;
    GetParameter("Gaussian SD", gaussianSD);

    double reductionRate = 0.8;
    GetParameter("Decimation rate", reductionRate);

    auto filter = ManualSegmentationToSurfaceFilter::New();
    filter->SetInput(binaryImage);
    filter->SetThreshold(0.5);
    filter->SetUseGaussianImageSmooth(smooth);
    filter->SetSmooth(smooth);
    filter->SetMedianFilter3D(applyMedian);

    if (smooth)
    {
      filter->InterpolationOn();
      filter->SetGaussianStandardDeviation(gaussianSD);
    }

    if (applyMedian)
      filter->SetMedianKernelSize(medianKernelSize, medianKernelSize, medianKernelSize);

    // Fix to avoid VTK warnings (see T5390)
    if (binaryImage->GetDimension() > 3)
      decimateMesh = false;

    if (decimateMesh)
    {
      filter->SetDecimate(ImageToSurfaceFilter::QuadricDecimation);
      filter->SetTargetReduction(reductionRate);
    }
    else
    {
      filter->SetDecimate(ImageToSurfaceFilter::NoDecimation);
    }

    filter->UpdateLargestPossibleRegion();

    auto surface = filter->GetOutput();
    auto polyData = surface->GetVtkPolyData();

    if (nullptr == polyData)
      throw std::logic_error("Could not create polygon model");

    polyData->SetVerts(nullptr);
    polyData->SetLines(nullptr);

    if (smooth || applyMedian || decimateMesh)
    {
      auto normals = vtkSmartPointer<vtkPolyDataNormals>::New();

      normals->AutoOrientNormalsOn();
      normals->FlipNormalsOff();
      normals->SetInputData(polyData);

      normals->Update();

      surface->SetVtkPolyData(normals->GetOutput());
    }
    else
    {
      surface->SetVtkPolyData(polyData);
    }

    return surface;
  }
}
