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

#include "smtk/session/aeva/Session.h"
#include "smtk/session/aeva/operators/LinearToQuadratic.h"

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

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

#include "smtk/operation/MarkGeometry.h"

#include "vtkCellArray.h"
#include "vtkCellData.h"
#include "vtkCellType.h"
#include "vtkDoubleArray.h"
#include "vtkExtractCells.h"
#include "vtkLinearToQuadraticCellsFilter.h"
#include "vtkMath.h"
#include "vtkPointData.h"
#include "vtkPolyData.h"
#include "vtkUnstructuredGrid.h"
#include "vtkUnstructuredGridGeometryFilter.h"
#include "vtkVector.h"
#include "vtkmCleanGrid.h"

#include <algorithm>

#include "smtk/session/aeva/LinearToQuadratic_xml.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
{

void promoteSideSets(const smtk::session::aeva::Resource::Ptr& resource,
  const smtk::session::aeva::Session::Ptr& session,
  const attribute::ComponentItem::Ptr& created,
  smtk::operation::MarkGeometry& geomMarker,
  const std::set<smtk::model::Entity*>& sideSets,
  const vtkSmartPointer<vtkUnstructuredGrid>& geom)
{
  for (auto* sideSet : sideSets)
  {
    auto* sideData = vtkPointSet::SafeDownCast(Operation::storage(sideSet));
    std::map<vtkIdType, int> cellCount;

    for (int i = 0; i < sideData->GetNumberOfPoints(); ++i)
    {
      auto pedigreeLocalPointId = geom->FindPoint(sideData->GetPoint(i));
      if (pedigreeLocalPointId < 0)
      {
        smtkErrorMacro(smtk::io::Logger::instance(), "No corresponding point found.");
      }

      vtkNew<vtkIdList> cellIds;
      geom->GetPointCells(pedigreeLocalPointId, cellIds);
      for (int j = 0; j < cellIds->GetNumberOfIds(); ++j)
      {
        auto cellId = cellIds->GetId(j);
        if (cellCount.find(cellId) != cellCount.end())
        {
          cellCount[cellId] = cellCount[cellId] + 1;
        }
        else
        {
          cellCount[cellId] = 1;
        }
      }
    }

    // Assume same cell type in the original geometry
    int nPointsInLinearCell = 0;
    switch (geom->GetCellType(0))
    {
      case VTK_QUADRATIC_EDGE:
        nPointsInLinearCell = 2;
        break;
      case VTK_QUADRATIC_TRIANGLE:
        nPointsInLinearCell = 3;
        break;
      case VTK_QUADRATIC_QUAD:
      case VTK_QUADRATIC_TETRA:
        nPointsInLinearCell = 4;
        break;
      case VTK_QUADRATIC_PYRAMID:
        nPointsInLinearCell = 5;
        break;
      case VTK_QUADRATIC_WEDGE:
        nPointsInLinearCell = 6;
        break;
      case VTK_QUADRATIC_HEXAHEDRON:
        nPointsInLinearCell = 8;
        break;
      default:
        smtkErrorMacro(smtk::io::Logger::instance(),
          "Found unsupported element type (" << geom->GetCellType(0) << ").");
        break;
    }

    vtkNew<vtkIdList> cellIds;
    cellIds->Allocate(cellCount.size());
    for (auto cellId : cellCount)
    {
      if (cellId.second == nPointsInLinearCell)
      {
        cellIds->InsertNextId(cellId.first);
      }
    }

    vtkNew<vtkExtractCells> extractCells;
    extractCells->SetCellList(cellIds);
    extractCells->SetInputData(geom);
    extractCells->Update();
    vtkNew<vtkUnstructuredGrid> newSideData;
    newSideData->ShallowCopy(extractCells->GetOutput());

    if (sideSet->isVertex())
    {
      auto npts = newSideData->GetNumberOfPoints();
      newSideData->GetCellData()->RemoveArray(
        newSideData->GetCellData()->GetGlobalIds()->GetName());
      newSideData->GetCellData()->SetNumberOfTuples(npts);
      vtkNew<vtkCellArray> cellArray;
      cellArray->SetNumberOfCells(npts);
      for (int i = 0; i < npts; ++i)
      {
        vtkIdType pts[1] = { i };
        cellArray->InsertNextCell(1, pts);
      }
      newSideData->SetCells(VTK_VERTEX, cellArray);
      auto newData = resource->addVertex();
      smtk::model::Vertex(sideSet->shared_from_this()).owningModel().addCell(newData);
      std::string name = sideSet->name();
      newData.setName(name.empty() ? "promoted" : (name + " (quadratic)"));
      auto newcomp = newData.entityRecord();
      session->addStorage(newcomp->id(), newSideData);
      session->setPrimary(*newcomp, false);
      created->appendValue(newcomp);
      geomMarker.markModified(newcomp);
    }
    else if (sideSet->isFace())
    {
      auto newData = resource->addFace();
      smtk::model::Face(sideSet->shared_from_this()).owningModel().addCell(newData);
      std::string name = sideSet->name();
      newData.setName(name.empty() ? "promoted" : (name + " (quadratic)"));
      auto newcomp = newData.entityRecord();
      session->addStorage(newcomp->id(), newSideData);
      session->setPrimary(*newcomp, false);
      created->appendValue(newcomp);
      geomMarker.markModified(newcomp);
    }
    else if (sideSet->isVolume())
    {
      auto newData = resource->addVolume();
      smtk::model::Volume(sideSet->shared_from_this()).owningModel().addCell(newData);
      std::string name = sideSet->name();
      newData.setName(name.empty() ? "promoted" : (name + " (quadratic)"));
      auto newcomp = newData.entityRecord();
      session->addStorage(newcomp->id(), newSideData);
      session->setPrimary(*newcomp, false);
      created->appendValue(newcomp);
      geomMarker.markModified(newcomp);
    }
    else
    {
      smtkErrorMacro(smtk::io::Logger::instance(), "Unsupported type of side set.");
    }
  }
}

bool LinearToQuadratic::ableToOperate()
{
  if (!this->Superclass::ableToOperate())
  {
    return false;
  }

  smtk::session::aeva::Resource::Ptr resource;
  smtk::session::aeva::Session::Ptr session;
  if (!this->fetchResourceAndSession(resource, session))
  {
    return false;
  }

  const std::vector<vtkIdType> linearElements = {
    VTK_LINE, VTK_TRIANGLE, VTK_QUAD, VTK_TETRA, VTK_HEXAHEDRON, VTK_WEDGE, VTK_PYRAMID
  };
  auto isLinearElement = [&linearElements](const vtkIdType& type) {
    return std::find(linearElements.begin(), linearElements.end(), type) != linearElements.end();
  };

  auto input = this->parameters()->associations()->valueAs<smtk::model::Entity>();
  if (!session->isPrimary(*input))
  {
    smtkErrorMacro(smtk::io::Logger::instance(), "Selection is not a primary geometry");
    return false;
  }

  auto data = Operation::storage(input);
  if (!data)
  {
    return false;
  }

  vtkSmartPointer<vtkPolyData> linPD(vtkPolyData::SafeDownCast(data));
  vtkSmartPointer<vtkUnstructuredGrid> linUG(vtkUnstructuredGrid::SafeDownCast(data));
  if (linUG)
  {
    if (linUG->GetNumberOfCells() == 0 || !isLinearElement(linUG->GetCellType(0)))
    {
      return false;
    }
  }
  else if (linPD)
  {
    if (linPD->GetNumberOfCells() == 0 || !isLinearElement(linPD->GetCellType(0)))
    {
      return false;
    }
  }
  else
  {
    return false;
  }
  return true;
}

LinearToQuadratic::Result LinearToQuadratic::operateInternal()
{
  smtk::session::aeva::Resource::Ptr resource;
  smtk::session::aeva::Session::Ptr session;
  auto result = this->createResult(smtk::operation::Operation::Outcome::FAILED);
  this->prepareResourceAndSession(result, resource, session);

  long pointIdOffset = 0;
  long maxPointId = 0;
  long cellIdOffset = 0;
  long maxCellId = 0;
  if (resource->properties().contains<long>("global point id offset"))
  {
    pointIdOffset = resource->properties().at<long>("global point id offset");
    // smtkInfoMacro(this->log(), "Found pt offset " << pointIdOffset );
  }
  if (resource->properties().contains<long>("global cell id offset"))
  {
    cellIdOffset = resource->properties().at<long>("global cell id offset");
    // smtkInfoMacro(this->log(), "Found cell offset " << cellIdOffset );
  }

  auto input = this->parameters()->associations()->valueAs<smtk::model::Entity>();
  auto data = Operation::storage(input);

  auto created = result->findComponent("created");
  smtk::operation::MarkGeometry geomMarker(resource);

  // Convert to unstructured grid for promotion
  vtkSmartPointer<vtkUnstructuredGrid> linUG;
  smtk::model::EntityPtr inputFace = nullptr;
  if (input->isFace())
  {
    // Check for parent volume
    bool hasVolume = false;
    auto volumes = input->relations();
    for (auto& volume : volumes)
    {
      auto volumePtr = input->resource()->find(volume)->as<smtk::model::Entity>();
      if (volumePtr->isVolume() && session->isPrimary(*volumePtr))
      {
        // Switch input to the primary volume
        // NOTE: only the first volume found is going to be promoted.
        inputFace = input;
        input = volumePtr;
        data = Operation::storage(volumePtr);
        linUG = vtkUnstructuredGrid::SafeDownCast(data);
        hasVolume = true;
        break;
      }
    }
    if (!hasVolume)
    {
      vtkSmartPointer<vtkPolyData> linPD = vtkPolyData::SafeDownCast(data);
      if (linPD)
      {
        vtkNew<vtkmCleanGrid> cleanGrid;
        cleanGrid->SetInputDataObject(linPD);
        cleanGrid->Update();
        linUG = cleanGrid->GetOutput();
      }
      else
      {
        linUG = vtkUnstructuredGrid::SafeDownCast(data);
      }
    }
  }
  else if (input->isVolume())
  {
    linUG = vtkUnstructuredGrid::SafeDownCast(data);
    auto faces = input->relations();
    for (auto& face : faces)
    {
      auto facePtr = input->resource()->find(face)->as<smtk::model::Entity>();
      if (facePtr->isFace() && session->isPrimary(*facePtr))
      {
        inputFace = facePtr;
      }
    }
  }
  else
  {
    smtkErrorMacro(smtk::io::Logger::instance(), "Unsupported primary geometry type");
  }

  // Quadratic promotion
  vtkNew<vtkLinearToQuadraticCellsFilter> promote;
  promote->SetInputDataObject(linUG);
  promote->Update();
  vtkNew<vtkUnstructuredGrid> geom;
  geom->ShallowCopy(promote->GetOutput());

  // Update the ids and offsets
  Session::offsetGlobalIds(geom, pointIdOffset, cellIdOffset, maxPointId, maxCellId);
  pointIdOffset = maxPointId + 1;
  cellIdOffset = maxCellId + 1;

  std::vector<vtkIdType> volumeType = { VTK_TETRA, VTK_HEXAHEDRON, VTK_WEDGE, VTK_PYRAMID };
  bool isVolume =
    std::find(volumeType.begin(), volumeType.end(), linUG->GetCellType(0)) != volumeType.end();

  // Add storage
  if (isVolume)
  {
    // Add volume
    auto volume = resource->addVolume();
    std::string name = input->name();
    volume.setName(name.empty() ? "promoted" : (name + " (quadratic)"));
    smtk::model::Volume(input).owningModel().addCell(volume);
    auto vcomp = volume.entityRecord();
    session->addStorage(vcomp->id(), geom);
    session->setPrimary(*vcomp, true);
    created->appendValue(vcomp);
    geomMarker.markModified(vcomp);

    // Add side sets
    geom->BuildLinks();
    std::set<smtk::model::Entity*> sideSets;
    if (session->sideSets(input.get(), sideSets))
    {
      // Remove node sets as they will be promoted from the boundary face
      // TODO: Remove this after adding function for selecting volume nodes:
      std::set<smtk::model::Entity*> newSideSets;
      for (auto* sideSet : sideSets)
      {
        if (!sideSet->isVertex())
        {
          newSideSets.insert(sideSet);
        }
      }
      promoteSideSets(resource, session, created, geomMarker, newSideSets, geom);
    }

    // Generate volume boundary
    vtkNew<vtkUnstructuredGridGeometryFilter> extractSurface;
    extractSurface->PassThroughPointIdsOn();
    extractSurface->SetInputDataObject(geom);
    extractSurface->Update();
    vtkNew<vtkUnstructuredGrid> surf;
    surf->ShallowCopy(extractSurface->GetOutput());

    // Pass through global point ids to the extracted boundary surface
    auto* surfPedigreeIds = surf->GetPointData()->GetGlobalIds("vtkOriginalPointIds");
    vtkNew<vtkIdList> surfPedigreeIdList;
    surfPedigreeIdList->SetNumberOfIds(surf->GetNumberOfPoints());
    for (int i = 0; i < surf->GetNumberOfPoints(); ++i)
    {
      surfPedigreeIdList->SetId(i, surfPedigreeIds->GetTuple1(i));
    }
    vtkNew<vtkIntArray> surfGlobalIds;
    surfGlobalIds->SetNumberOfValues(surf->GetNumberOfPoints());
    geom->GetPointData()->GetGlobalIds()->GetTuples(surfPedigreeIdList, surfGlobalIds);
    surfGlobalIds->SetName("GlobalPointIds");
    surf->GetPointData()->SetGlobalIds(surfGlobalIds);

    //Update the cell ID offset, but leave global point IDs alone - they pass through.
    Session::offsetGlobalIds(surf, -1, cellIdOffset, maxPointId, maxCellId);
    pointIdOffset = maxPointId + 1;
    cellIdOffset = maxCellId + 1;

    // Add boundary
    auto face = resource->addFace();
    std::string prefix = "boundary of ";
    face.setName(prefix + volume.name());
    smtk::model::EntityRef(volume).addRawRelation(face);
    smtk::model::EntityRef(face).addRawRelation(volume);
    auto fcomp = face.entityRecord();
    session->addStorage(fcomp->id(), surf);
    session->setPrimary(*fcomp, true);
    created->appendValue(fcomp);
    geomMarker.markModified(fcomp);

    // Add side sets
    surf->BuildLinks();
    std::set<smtk::model::Entity*> faceSideSets;
    if (inputFace && session->sideSets(inputFace.get(), faceSideSets))
    {
      promoteSideSets(resource, session, created, geomMarker, faceSideSets, surf);
    }
  }
  else
  {
    // Add surface
    auto face = resource->addFace();
    std::string name = input->name();
    face.setName(name.empty() ? "promoted" : (name + " (quadratic)"));
    smtk::model::Face(input).owningModel().addCell(face);
    auto fcomp = face.entityRecord();
    session->addStorage(fcomp->id(), geom);
    session->setPrimary(*fcomp, true);
    created->appendValue(fcomp);
    geomMarker.markModified(fcomp);

    // Add side sets
    geom->BuildLinks();
    std::set<smtk::model::Entity*> sideSets;
    if (session->sideSets(input.get(), sideSets))
    {
      promoteSideSets(resource, session, created, geomMarker, sideSets, geom);
    }
  }

  resource->properties().get<long>()["global point id offset"] = pointIdOffset;
  resource->properties().get<long>()["global cell id offset"] = cellIdOffset;

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

const char* LinearToQuadratic::xmlDescription() const
{
  return LinearToQuadratic_xml;
}

}
}
}
