//=============================================================================
// Copyright (c) Kitware, Inc.
// All rights reserved.
// See LICENSE.txt for details.
//
// This software is distributed WITHOUT ANY WARRANTY; without even
// the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE.  See the above copyright notice for more information.
//=============================================================================

#include "smtk/session/aeva/operators/ImprintGeometry.h"
#include "smtk/session/aeva/operators/ImprintGeometry_xml.h"

#include "smtk/attribute/Attribute.h"
#include "smtk/attribute/DoubleItem.h"
#include "smtk/attribute/IntItem.h"
#include "smtk/attribute/StringItem.h"
#include "smtk/attribute/VoidItem.h"
#include "smtk/io/Logger.h"
#include "smtk/model/Volume.h"
#include "smtk/operation/MarkGeometry.h"
#include "vtk/aeva/ext/vtkImageNarrowBand.h"

#include "vtkImageData.h"
#include "vtkMultiThreader.h"
#include "vtkPointData.h"
#include "vtkPolyData.h"
#include "vtkShepardMethod.h"
#include "vtkUnstructuredGrid.h"

#include "vtkNIFTIImageWriter.h"

// On Windows MSVC 2015+, something is included that defines
// a macro named ERROR to be 0. This causes smtkErrorMacro()
// to expand into garbage (because smtk::io::Logger::ERROR
// gets expanded to smtk::io::Logger::0).
#ifdef ERROR
#undef ERROR
#endif

namespace smtk
{
namespace session
{
namespace aeva
{

ImprintGeometry::Result ImprintGeometry::operateInternal()
{
  smtk::attribute::ReferenceItemPtr assocs = this->parameters()->associations();

  smtk::operation::Operation::Result result =
    this->createResult(smtk::operation::Operation::Outcome::FAILED);

  if (assocs->numberOfValues() < 2)
  {
    smtkErrorMacro(this->log(),
      "ImprintGeometry operation requires 2 as associations, an image and mesh association");
    return result;
  }

  smtk::model::Model parentModel1 =
    smtk::model::EntityRef(assocs->valueAs<smtk::model::Entity>(0)).owningModel();
  smtk::model::Model parentModel2 =
    smtk::model::EntityRef(assocs->valueAs<smtk::model::Entity>(1)).owningModel();
  smtk::model::Model parentModel =
    parentModel1; // Will take the value of 1 or 2 depending on which contains vtkImageData

  // Identify which association is an image and which is vtkUnstructuredGrid or vtkPolyData
  vtkSmartPointer<vtkDataSet> data1 =
    vtkDataSet::SafeDownCast(Operation::getData(assocs->value(0)));
  vtkSmartPointer<vtkDataSet> data2 =
    vtkDataSet::SafeDownCast(Operation::getData(assocs->value(1)));
  if (!data1 || !data2)
  {
    smtkErrorMacro(this->log(), "ImprintGeometry: Unable to retrieve geometry for input data");
    return result;
  }

  const std::string class1 = data1->GetClassName();
  const std::string class2 = data2->GetClassName();

  vtkSmartPointer<vtkImageData> input1;
  vtkSmartPointer<vtkDataSet> input2;
  if (class1 == "vtkImageData" && (class2 == "vtkPolyData" || class2 == "vtkUnstructuredGrid"))
  {
    input1 = vtkImageData::SafeDownCast(data1);
    input2 = data2;
  }
  else if (class2 == "vtkImageData" && (class1 == "vtkPolyData" || class1 == "vtkUnstructuredGrid"))
  {
    input1 = vtkImageData::SafeDownCast(data2);
    parentModel = parentModel2; // If data2 is the image, then switch the parent model
    input2 = data1;
  }
  else
  {
    smtkErrorMacro(this->log(),
      "ImprintGeometry: wrong data types, must use imagedata + polydata/unstructuredgrid");
    return result;
  }

  // Create a volume within the parentModel
  auto modelComp = parentModel.component();
  modelComp->properties().get<std::string>()["aeva_datatype"] = "image";
  auto volume = parentModel.resource()->addVolume();

  auto imprintMethod = this->parameters()->findString("imprintMethod")->value(0);
  if (imprintMethod == "contact" || imprintMethod == "distance")
  {
    vtkNew<vtkImageNarrowBand> narrowBand;
    // we want multi-threading, but ParaView sets to 1 by default. Override
    int threadMax = vtkMultiThreader::GetGlobalMaximumNumberOfThreads();
    vtkMultiThreader::SetGlobalMaximumNumberOfThreads(
      vtkMultiThreader::GetGlobalStaticMaximumNumberOfThreads());

    narrowBand->SetReferenceImage(input1);
    narrowBand->SetInputData(0, input2);
    narrowBand->SetBandWidth(this->parameters()->findDouble("bandwidth")->value(0));
    narrowBand->SetUseBinary(imprintMethod == "contact");
    narrowBand->Update();

    // always sets either "Band" or "Distance" as active scalar array
    if (!narrowBand->GetOutput()->GetPointData()->GetScalars())
    {
      smtkErrorMacro(this->log(), "ImprintGeometry: vtkImageNarrowBand failure.");
      return result;
    }

    parentModel.resource()->as<smtk::session::aeva::Resource>()->session()->addStorage(
      volume.entity(), narrowBand->GetOutput());
    volume.setName(narrowBand->GetOutput()->GetPointData()->GetScalars()->GetName());

    vtkMultiThreader::SetGlobalMaximumNumberOfThreads(threadMax);
  }
  else if (imprintMethod == "scalar")
  {
    // DEBUG
    std::cout << "DBG Shepard " << input2->GetPointData()->GetNumberOfArrays() << std::endl;
    auto scalarName = this->parameters()->findString("scalarName");
    if (scalarName && scalarName->isEnabled())
    {
      std::string activeName = scalarName->value();
      if (!activeName.empty())
      {
        input2->GetPointData()->SetActiveScalars(activeName.c_str());
      }
    }
    if (!input2->GetPointData()->GetScalars())
    {
      smtkErrorMacro(this->log(), "ImprintGeometry: vtkShepardMethod requires active scalars");
      return result;
    }
    double maxDist = 0.0;
    auto maxDistItem = this->parameters()->findDouble("maxInfluence");
    if (maxDistItem && maxDistItem->isEnabled())
    {
      maxDist = maxDistItem->value();
    }
    // vtkShepardMethod or potentially vtkGaussianSplatter
    int dim[3];
    double bounds[6];
    input1->GetDimensions(dim);
    input1->GetBounds(bounds);

    vtkNew<vtkShepardMethod> shepard;
    shepard->SetInputData(0, input2);
    shepard->SetSampleDimensions(dim);
    shepard->SetModelBounds(bounds);
    // the max dist input is relative to the max of the model bounds, so scale the user's input.
    if (maxDist > 0.0)
    {
      double scaledMax = std::min(1.0,
        maxDist /
          std::max(bounds[1] - bounds[0], std::max(bounds[3] - bounds[2], bounds[5] - bounds[4])));
      if (scaledMax > 0.0)
      {
        shepard->SetMaximumDistance(scaledMax);
      }
    }
    shepard->Update();
    if (!shepard->GetOutput()->GetPointData()->GetScalars())
    {
      smtkErrorMacro(this->log(), "ImprintGeometry: vtkShepardMethod failure.");
      return result;
    }

    parentModel.resource()->as<smtk::session::aeva::Resource>()->session()->addStorage(
      volume.entity(), shepard->GetOutput());
    volume.setName(shepard->GetOutput()->GetPointData()->GetScalars()->GetName());
  }
  parentModel.addCell(volume);
  auto createdItems = result->findComponent("created");
  createdItems->appendValue(volume.component());
  operation::MarkGeometry(parentModel.resource()).markModified(volume.component());

  result->findInt("outcome")->setValue(static_cast<int>(ImprintGeometry::Outcome::SUCCEEDED));
  return result;
}

const char* ImprintGeometry::xmlDescription() const
{
  return ImprintGeometry_xml;
}

}
}
}
