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

#include "smtk/session/aeva/Export_xml.h"
#include "smtk/session/aeva/Resource.h"
#include "smtk/session/aeva/Session.h"
#include "vtk/aeva/ext/vtkMedWriter.h"

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

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

#include "smtk/resource/Manager.h"

#include "smtk/common/Paths.h"

#include "vtkCell3D.h"
#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkSmartPointer.h"
#include "vtkUnstructuredGrid.h"
#include "vtkXMLImageDataWriter.h"

#ifdef ERROR
#undef ERROR
#endif

using ResourceArray = std::vector<smtk::session::aeva::Resource::Ptr>;

namespace smtk
{
namespace session
{
namespace aeva
{

Export::Result Export::operateInternal()
{
  smtk::attribute::FileItem::Ptr filenameItem = this->parameters()->findFile("filename");

  std::string filename = filenameItem->value();

  // Infer file type from name
  std::string ext = smtk::common::Paths::extension(filename);
  // Downcase the extension
  std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
  if (ext == ".med")
  {
    return this->exportMedMesh();
  }

  // Internal storage is a VTK image, default to that:
  return this->exportVTKImage();
}

Export::Result Export::exportVTKImage()
{
  auto inputItem = this->parameters()->associations();
  auto inputValue = inputItem->valueAs<smtk::model::Entity>(0);
  smtk::session::aeva::Resource::Ptr resource;
  if (inputValue)
  {
    resource = std::static_pointer_cast<smtk::session::aeva::Resource>(inputValue->resource());
  }
  if (!resource)
  {
    smtkErrorMacro(this->log(), "No resources to export.");
    return this->createResult(smtk::operation::Operation::Outcome::FAILED);
  }

  smtk::common::UUID imageId;
  if (inputValue->isModel())
  {
    // retrieve the first cell, containing the volume image to save.
    smtk::model::Model dataset = smtk::model::Model(resource, inputValue->id());
    if (!dataset.cells().empty())
    {
      imageId = dataset.cells()[0].entity();
    }
  }
  else if (inputValue->isCellEntity())
  {
    imageId = inputValue->id();
  }

  auto const& session = resource->session();
  auto* image = vtkImageData::SafeDownCast(session->findStorage(imageId));
  if (!image)
  {
    smtkErrorMacro(this->log(), "Cannot retrieve image from model.");
    return this->createResult(smtk::operation::Operation::Outcome::FAILED);
  }

  vtkNew<vtkXMLImageDataWriter> writer;
  // force file extension to match what we're writing.
  std::string filename = smtk::common::Paths::replaceExtension(
    this->parameters()->findFile("filename")->value(), ".vti");

  writer->SetFileName(filename.c_str());
  writer->SetInputData(image);
  writer->Write();
  auto result = this->createResult(Export::Outcome::SUCCEEDED);
  return result;
}

Export::Result Export::exportMedMesh()
{
  auto inputItem = this->parameters()->associations();

  vtkNew<vtkMedWriter> writer;
  // force file extension to match what we're writing.
  std::string filename = smtk::common::Paths::replaceExtension(
    this->parameters()->findFile("filename")->value(), ".med");
  writer->SetFileName(filename);

  for (int k = 0; k < inputItem->numberOfValues(); k++)
  {
    auto inputValue = inputItem->valueAs<smtk::model::Entity>(k);
    smtk::session::aeva::Resource::Ptr resource = nullptr;
    smtk::session::aeva::Session::Ptr session = nullptr;
    m_result = this->createResult(Export::Outcome::FAILED);

    this->prepareResourceAndSession(m_result, resource, session, false);
    if (!resource)
    {
      smtkErrorMacro(this->log(), "No resources to export.");
      return m_result;
    }

    // organize data from the resource in a multiblock as the writer expects.
    vtkSmartPointer<vtkMultiBlockDataSet> writerInput = vtkMultiBlockDataSet::New();
    vtkSmartPointer<vtkMultiBlockDataSet> geomInput = vtkMultiBlockDataSet::New();
    writerInput->SetBlock(0, geomInput);
    vtkSmartPointer<vtkMultiBlockDataSet> groupInput = vtkMultiBlockDataSet::New();
    writerInput->SetBlock(1, groupInput);

    // TODO: Currently available med files start with a single surface,
    // and an optional encoded tetra volume. This may not be true in the future.
    // We always create a volume on import, even if it's empty.
    // Look for the volume and surface, add to the geom MB.
    // Side-sets, or groups, are marked with a property flag and added to the group MB.
    smtk::common::UUID volId;
    if (!inputValue->isModel())
    {
      smtkErrorMacro(this->log(), "Expected model as input to export.");
      return m_result;
    }
    smtk::model::Model dataset = smtk::model::Model(resource, inputValue->id());
    vtkSmartPointer<vtkUnstructuredGrid> volume;
    std::string volumeName;
    vtkSmartPointer<vtkUnstructuredGrid> surface;
    std::string surfaceName;
    for (auto const& cellRef : dataset.cells())
    {
      auto const& cellId = cellRef.entity();
      auto* unstructuredGrid = vtkUnstructuredGrid::SafeDownCast(session->findStorage(cellId));
      if (unstructuredGrid && Session::isSideSet(*cellRef.component()))
      {
        // add to the group multi-block
        auto i = groupInput->GetNumberOfBlocks();
        groupInput->SetBlock(i, unstructuredGrid);
        groupInput->GetMetaData(i)->Set(vtkCompositeDataSet::NAME(), cellRef.name());
      }
      else if (!volume && cellRef.dimension() == 3 &&
        (!unstructuredGrid ||
          (unstructuredGrid->GetNumberOfCells() > 0 &&
            vtkCell3D::SafeDownCast(unstructuredGrid->GetCell(0)) != NULL)))
      {
        // unstructuredGrid can be empty, if the surface encloses a rigid volume with no tetra mesh.
        volume = unstructuredGrid;
        volumeName = cellRef.name();
        // retrieve the surface from the volume relations
        for (auto const& relation : cellRef.relations())
        {
          if (relation.dimension() == 2)
          {
            unstructuredGrid =
              vtkUnstructuredGrid::SafeDownCast(session->findStorage(relation.entity()));
            if (unstructuredGrid && unstructuredGrid->GetCellType(0) == VTK_TRIANGLE)
            {
              surface = unstructuredGrid;
              surfaceName = relation.name();
            }
          }
        }
      }
      else if (unstructuredGrid && !surface && cellRef.dimension() == 2 &&
        (unstructuredGrid->GetNumberOfCells() > 0 &&
          vtkCell3D::SafeDownCast(unstructuredGrid->GetCell(0)) == NULL))
      {
        surface = unstructuredGrid;
        surfaceName = cellRef.name();
      }
    }

    if (!surface)
    {
      smtkErrorMacro(this->log(), "Cannot retrieve surface from model.");
      return m_result;
    }
    unsigned int block = 0;
    if (volume)
    {
      geomInput->SetBlock(block, volume);
      geomInput->GetMetaData(block)->Set(vtkCompositeDataSet::NAME(), volumeName);
      ++block;
    }
    geomInput->SetBlock(block, surface);
    geomInput->GetMetaData(block)->Set(vtkCompositeDataSet::NAME(), surfaceName);
    writerInput->GetInformation()->Set(vtkMultiBlockDataSet::NAME(), inputValue->name());
    writer->AddInputDataObject(0, writerInput);
  }

  writer->Update();

  m_result->findInt("outcome")->setValue(0, static_cast<int>(Export::Outcome::SUCCEEDED));
  return m_result;
}

const char* Export::xmlDescription() const
{
  return Export_xml;
}

bool exportResource(const smtk::resource::ResourcePtr& resource)
{
  Export::Ptr exportResource = Export::create();
  exportResource->parameters()->findResource("resource")->setValue(resource);
  Export::Result result = exportResource->operate();
  return (result->findInt("outcome")->value() == static_cast<int>(Export::Outcome::SUCCEEDED));
}
}
}
}
