//=========================================================================
//  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/Registrar.h"

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

#include "smtk/attribute/Attribute.h"
#include "smtk/attribute/FileItem.h"
#include "smtk/attribute/IntItem.h"
#include "smtk/attribute/StringItem.h"

#include "smtk/model/Face.h"
#include "smtk/model/Volume.h"

#include <vtkDoubleArray.h>
#include <vtkPointData.h>
#include <vtkPolyData.h>

namespace
{
std::string dataRoot = AEVA_DATA_DIR;

smtk::model::Entity::Ptr createTestSurface(smtk::operation::Manager::Ptr const& operationManager)
{
  // Create an import operation
  smtk::session::aeva::Import::Ptr importOp =
    operationManager->create<smtk::session::aeva::Import>();
  if (!importOp)
  {
    std::cerr << "No import operation\n";
    return nullptr;
  }

  // Set the file path
  std::string importFilePath(dataRoot);
  importFilePath += "/vtk/oks003_TBB_AGS.vtk";
  importOp->parameters()->findFile("filename")->setValue(importFilePath);

  // Test the ability to operate
  if (!importOp->ableToOperate())
  {
    std::cerr << "Import operation unable to operate\n";
    return nullptr;
  }

  // Execute the operation
  smtk::operation::Operation::Result importOpResult = importOp->operate();

  // Test for success
  if (importOpResult->findInt("outcome")->value() !=
    static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
  {
    std::cerr << "Import operation failed\n";
    return nullptr;
  }

  // Retrieve the resulting model
  smtk::attribute::ComponentItemPtr componentItem = importOpResult->findComponent("created");

  // Access the generated model
  auto model = std::dynamic_pointer_cast<smtk::model::Entity>(componentItem->value());

  // Test model validity
  if (!model->referenceAs<smtk::model::Model>().isValid())
  {
    std::cerr << "Imported model is invalid\n";
    return nullptr;
  }

  return model;
}

// Creates test volume in resource
smtk::model::Volume createTestVolume(smtk::session::aeva::Resource::Ptr const& resource,
  double* bounds,
  int* dim)
{
  resource->setName("testImage");

  vtkNew<vtkImageData> image;
  image->SetDimensions(dim);
  double spacing[3] = { (bounds[1] - bounds[0]) / static_cast<double>(dim[0]),
    (bounds[3] - bounds[2]) / static_cast<double>(dim[1]),
    (bounds[5] - bounds[4]) / static_cast<double>(dim[2]) };
  image->SetSpacing(spacing);
  image->SetOrigin(bounds[0], bounds[2], bounds[4]);
  int extent[6]{ 0, dim[0] - 1, 0, dim[1] - 1, 0, dim[2] - 1 };
  image->SetExtent(extent);
  image->AllocateScalars(VTK_UNSIGNED_CHAR, 1);
  image->GetPointData()->GetScalars()->Fill(0);

  auto model = resource->addModel(3, 3, "testImage");
  auto modelComp = model.component();
  modelComp->properties().get<std::string>()["aeva_datatype"] = "image";
  auto volume = resource->addVolume();
  model.addCell(volume);

  resource->session()->addStorage(volume.entity(), image);
  if (image->GetPointData()->GetScalars())
  {
    volume.setName(image->GetPointData()->GetScalars()->GetName());
  }
  return volume;
}
}

int TestImprintGeometryOp(int argc, char* argv[])
{
  (void)argc;
  (void)argv;

  // Create a resource manager
  smtk::resource::Manager::Ptr resourceManager = smtk::resource::Manager::create();

  // Register the aeva session to the resource manager
  {
    smtk::session::aeva::Registrar::registerTo(resourceManager);
  }

  // Create an operation manager
  smtk::operation::Manager::Ptr operationManager = smtk::operation::Manager::create();

  // Register the aeva session to the operation manager
  {
    smtk::session::aeva::Registrar::registerTo(operationManager);
  }

  // Register the resource manager to the operation manager (newly created
  // resources will be automatically registered to the resource manager).
  operationManager->registerResourceManager(resourceManager);

  smtk::model::Entity::Ptr surfModel = createTestSurface(operationManager);
  if (surfModel == nullptr)
  {
    return 1;
  }
  smtk::session::aeva::Session::Ptr session =
    std::static_pointer_cast<smtk::session::aeva::Resource>(surfModel->modelResource())->session();
  // Access all of the model's faces
  smtk::model::EntityRefs faces =
    surfModel->modelResource()->entitiesMatchingFlagsAs<smtk::model::EntityRefs>(smtk::model::FACE);
  if (faces.empty())
  {
    std::cerr << "No faces\n";
    return 1;
  }
  smtk::model::Face face = *faces.begin();
  vtkSmartPointer<vtkPolyData> faceData =
    vtkPolyData::SafeDownCast(session->findStorage(face.entity()));

  // Add a scalar field to test imprinting scalars.
  auto* points = faceData->GetPointData();
  auto numPts = faceData->GetNumberOfPoints();
  vtkSmartPointer<vtkDoubleArray> testScalars = vtkSmartPointer<vtkDoubleArray>::New();
  testScalars->SetName("ramp");
  testScalars->SetNumberOfComponents(1);
  testScalars->SetNumberOfValues(numPts);
  for (auto i = 0; i < numPts; ++i)
  {
    testScalars->SetValue(i, static_cast<double>(i) / numPts);
  }
  points->AddArray(testScalars);
  points->SetActiveScalars("ramp");

  int dims[3] = { 40, 40, 40 };
  smtk::session::aeva::Resource::Ptr resource =
    resourceManager->create<smtk::session::aeva::Resource>();
  resource->setSession(session);
  smtk::model::Volume volume = createTestVolume(resource, faceData->GetBounds(), dims);

  // Create an ImprintGeometry operation
  smtk::session::aeva::ImprintGeometry::Ptr imprintGeometry =
    operationManager->create<smtk::session::aeva::ImprintGeometry>();
  if (!imprintGeometry)
  {
    std::cerr << "No imprint geometry operation\n";
    return 1;
  }

  // Set the input face
  imprintGeometry->parameters()->associate(face.component());
  imprintGeometry->parameters()->associate(volume.component());

  // Test the ability to operate
  if (!imprintGeometry->ableToOperate())
  {
    std::cerr << "Imprint geometry operation unable to operate\n";
    return 1;
  }

  // Execute the operation
  smtk::operation::Operation::Result imprintGeometryResult = imprintGeometry->operate();

  // Test for success
  if (imprintGeometryResult->findInt("outcome")->value() !=
    static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
  {
    std::cerr << "Imprint geometry operation failed\n";
    return 1;
  }

  {
    // The operation would produce another volume on the resource associated with volume
    smtk::attribute::ComponentItemPtr resultsCompItem =
      imprintGeometryResult->findComponent("created");
    smtk::model::Volume vol =
      resultsCompItem->value(0)->as<smtk::model::Entity>()->referenceAs<smtk::model::Volume>();

    vtkSmartPointer<vtkImageData> resultImageData =
      vtkImageData::SafeDownCast(session->findStorage(vol.entity()));
    if (!resultImageData)
    {
      std::cerr << "Missing expected output data\n";
      return 1;
    }

    auto* resultPtr = static_cast<unsigned char*>(resultImageData->GetScalarPointer());
    int dim[3];
    resultImageData->GetDimensions(dim);
    int sum = 0;
    for (int i = 0; i < dim[0] * dim[1] * dim[2]; i++)
    {
      sum += static_cast<int>(resultPtr[i]);
    }

    if (sum != 12358)
    {
      std::cerr << "Invalid binary image sum " << sum << std::endl;
      return 1;
    }
  }

  // ======= next pass
  imprintGeometry->parameters()->findString("imprintMethod")->setValue("scalar");
  if (!imprintGeometry->ableToOperate())
  {
    std::cerr << "Imprint geometry operation 2 unable to operate\n";
    return 1;
  }

  // Execute the operation
  imprintGeometryResult = imprintGeometry->operate();

  // Test for success
  if (imprintGeometryResult->findInt("outcome")->value() !=
    static_cast<int>(smtk::operation::Operation::Outcome::SUCCEEDED))
  {
    std::cerr << "Imprint geometry operation 2 failed\n";
    return 1;
  }

  {
    // The operation would produce another volume on the resource associated with volume
    smtk::attribute::ComponentItemPtr resultsCompItem =
      imprintGeometryResult->findComponent("created");
    smtk::model::Volume vol =
      resultsCompItem->value(0)->as<smtk::model::Entity>()->referenceAs<smtk::model::Volume>();

    vtkSmartPointer<vtkImageData> resultImageData =
      vtkImageData::SafeDownCast(session->findStorage(vol.entity()));
    if (!resultImageData)
    {
      std::cerr << "Missing expected output data\n";
      return 1;
    }

    // shepard's method produces a float image.
    auto* resultPtr = static_cast<float*>(resultImageData->GetScalarPointer());
    int dim[3];
    resultImageData->GetDimensions(dim);
    int sum = 0;
    for (int i = 0; i < dim[0] * dim[1] * dim[2]; i++)
    {
      sum += (resultPtr[i] > 0.5 ? 1 : 0);
    }

    std::cout << "Shepard " << sum << " of " << dim[0] * dim[1] * dim[2] << std::endl;
    if (sum != 11564)
    {
      std::cerr << "Invalid ramp image count " << sum << std::endl;
      return 1;
    }
  }
  return 0;
}
