//=========================================================================
//=========================================================================
//  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/ExportFeb.h"
#include "smtk/session/aeva/ExportFeb_xml.h"

#include "smtk/session/aeva/Resource.h"
#include "smtk/session/aeva/Session.h"

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

#include "smtk/common/Paths.h"
#include "smtk/model/EntityRef.h"
#include "smtk/model/Face.h"
#include "smtk/model/Volume.h"

#include "vtkCellData.h"
#include "vtkIntArray.h"
#include "vtkPointData.h"
#include "vtkPolyData.h"
#include "vtkSmartPointer.h"
#include "vtk_pugixml.h"

#ifdef ERROR
#undef ERROR
#endif

using Offsets = std::unordered_map<smtk::common::UUID, int>;

namespace smtk
{
namespace session
{
namespace aeva
{

namespace
{

std::vector<std::string> getCellNames()
{
  std::vector<std::string> cellNames(VTK_NUMBER_OF_CELL_TYPES);
  cellNames[VTK_LINE] = "line2";
  cellNames[VTK_TRIANGLE] = "tri3";
  cellNames[VTK_QUAD] = "quad4";
  cellNames[VTK_TETRA] = "tet4";
  cellNames[VTK_HEXAHEDRON] = "hex8";
  cellNames[VTK_QUADRATIC_EDGE] = "line3";
  cellNames[VTK_QUADRATIC_TRIANGLE] = "tri6";
  cellNames[VTK_QUADRATIC_QUAD] = "quad9";

  return cellNames;
}

// from vtk cell type to cell name in febio
const std::vector<std::string> cellNames(getCellNames());

/*
 * \brief Get the range of global point/cell ids of an geometry
 * \return false if the backend storage is empty, otherwise true.
 */
template<bool PointId>
bool getRangeOfGlobalIds(const smtk::session::aeva::Session::Ptr& session,
  const smtk::common::UUID& id,
  int& minId,
  int& maxId)
{
  vtkSmartPointer<vtkPointSet> mesh(vtkPointSet::SafeDownCast(session->findStorage(id)));
  minId = 0;
  maxId = 0;
  if (!mesh)
  {
    return false;
  }
  vtkIntArray* ids;
  if (PointId)
  {
    ids = vtkIntArray::SafeDownCast(mesh->GetPointData()->GetGlobalIds());
  }
  else
  {
    ids = vtkIntArray::SafeDownCast(mesh->GetCellData()->GetGlobalIds());
  }
  int* range = ids->GetValueRange();
  minId = range[0];
  maxId = range[1];
  return true;
}

/*
 * \brief Add nodes of a mesh to xml node
 *
 * \param mesh the backend vtk mesh of a primary geometry
 * \param name name of the mesh
 * \param offset offset in node indices
 * \param xmlNode a xml node
 */
bool addNodes(vtkPointSet& mesh,
  const std::string& name,
  const int febOffset,
  pugi::xml_node& xmlNode)
{
  xmlNode.set_name("Nodes");
  xmlNode.append_attribute("name") = name.c_str();

  vtkIdType numPts = mesh.GetNumberOfPoints();
  vtkPoints* pnts = mesh.GetPoints();
  vtkIdType id;
  double xyz[3];

  for (vtkIdType ii = 0; ii < numPts; ++ii)
  {
    id = ii + febOffset + 1;
    pnts->GetPoint(ii, xyz);
    pugi::xml_node pnt = xmlNode.append_child(std::to_string(id).c_str());
    pnt.set_name("node");
    pnt.append_attribute("id") = std::to_string(id).c_str();
    std::string coords =
      " " + std::to_string(xyz[0]) + ", " + std::to_string(xyz[1]) + ", " + std::to_string(xyz[2]);
    pnt.append_child(pugi::node_pcdata).set_value(coords.c_str());
  }

  return true;
}

/*
 * \brief Add elements of a mesh to xml node
 *
 * \param mesh the backend vtk mesh of a primary geometry
 * \param name name of the mesh
 * \param pntOffset offset in point indices
 * \param cellOffset offset in element indices
 * \param xmlNode a xml node
 */
bool addElements(vtkPointSet& mesh,
  const std::string& name,
  const int pntOffset,
  const int cellOffset,
  pugi::xml_node& xmlNode)
{
  xmlNode.set_name("Elements");
  xmlNode.append_attribute("type") = cellNames[mesh.GetCellType(0)].c_str(); // element type
  xmlNode.append_attribute("mat") = "1";                                     // material id
  xmlNode.append_attribute("name") = name.c_str();

  for (vtkIdType ii = 0; ii < mesh.GetNumberOfCells(); ++ii)
  {
    vtkIdType id = ii + cellOffset + 1; // feb id
    auto* cell = mesh.GetCell(ii);
    pugi::xml_node pnt = xmlNode.append_child(std::to_string(id).c_str());
    pnt.set_name("elem");
    pnt.append_attribute("id") = std::to_string(id).c_str();

    // \todo: use setwidth
    std::string delimiter;
    std::string indices;
    for (vtkIdType jj = 0; jj < cell->GetNumberOfPoints(); ++jj)
    {
      delimiter = (jj == 0) ? " " : ", ";
      indices += delimiter + std::to_string(cell->GetPointId(jj) + pntOffset + 1);
    }
    pnt.append_child(pugi::node_pcdata).set_value(indices.c_str());
  }

  return true;
}

bool addNodeSet(smtk::model::Entity* ent,
  smtk::session::aeva::Session::Ptr session,
  smtk::io::Logger& log,
  const Offsets& pntOffset,
  pugi::xml_node& xmlNode)
{
  // get vtk mesh
  vtkSmartPointer<vtkPointSet> sideset(vtkPointSet::SafeDownCast(session->findStorage(ent->id())));
  if (!sideset)
  {
    smtkErrorMacro(log, "Could not convert to vtkPointSet.");
    return false;
  }
  vtkIntArray* sideSetPntIds = vtkIntArray::SafeDownCast(sideset->GetPointData()->GetGlobalIds());
  if (sideSetPntIds == nullptr)
  {
    smtkErrorMacro(log, "global point ids are missing");
    return false;
  }

  xmlNode.set_name("NodeSet");
  xmlNode.append_attribute("name") = ent->name().c_str();
  std::set<smtk::model::Entity*> primaries;
  session->primaryGeometries(ent, primaries);
  // track if a vertex cell is processed. By process, I mean its corresponding
  // ids of its points are found in a primary mesh
  std::vector<bool> found(sideset->GetNumberOfCells(), false);
  // track how many vertex cells are processed.
  int numProcessed = 0;

  // Process nodes that belong to a primary.
  auto processNodesInPrimary = [&](smtk::model::Entity* primary) {
    int minGlobalId;
    int maxGlobalId;
    auto foundRange = getRangeOfGlobalIds<true>(session, primary->id(), minGlobalId, maxGlobalId);
    if (!foundRange)
    {
      // continue;
      return;
    }
    int febOffset = pntOffset.at(primary->id());

    // In some cases some nodes are not referenced by the cells,
    // so loop over cells instead of points.
    for (vtkIdType ii = 0; ii < sideset->GetNumberOfCells(); ++ii)
    {
      auto* cell = sideset->GetCell(ii);
      int id = sideSetPntIds->GetValue(cell->GetPointId(0));
      if (found[ii] || id < minGlobalId || id > maxGlobalId)
      {
        continue;
      }
      id += febOffset - minGlobalId + 1;
      pugi::xml_node pnt = xmlNode.append_child(std::to_string(id).c_str());
      pnt.set_name("node");
      pnt.append_attribute("id") = std::to_string(id).c_str();
      found[ii] = true;
      ++numProcessed;
    }
  };

  for (auto* primary : primaries)
  {
    if (primary->isVolume())
    {
      processNodesInPrimary(primary);
    }
  }

  if (numProcessed == sideset->GetNumberOfCells())
  {
    return true;
  }

  // If we get here, it means one or more volume primaries are empty,
  // or there is no volume primaries at all. We have to search unprocessed
  // nodes in surface primaries.
  for (auto* primary : primaries)
  {
    if (primary->isFace() && pntOffset.find(primary->id()) != pntOffset.end())
    {
      processNodesInPrimary(primary);
    }
  }

  // To see if all points are found.
  if (numProcessed != sideset->GetNumberOfCells())
  {
    std::string msg = "side-set " + ent->name() + ": " + std::to_string(numProcessed) + " out of " +
      std::to_string(sideset->GetNumberOfCells()) + " found in primaries.";
    smtkErrorMacro(log, msg);
    return false;
  }

  return true;
}

bool addFaceSet(smtk::model::Entity* ent,
  const smtk::session::aeva::Session::Ptr& session,
  smtk::io::Logger& log,
  const Offsets& pntOffset,
  pugi::xml_node& xmlNode)
{
  // convert to vtk mesh
  vtkSmartPointer<vtkPointSet> sideset(vtkPointSet::SafeDownCast(session->findStorage(ent->id())));
  if (!sideset)
  {
    smtkErrorMacro(log, "Could not convert to vtkPointSet.");
    return false;
  }
  vtkIntArray* pntIds = vtkIntArray::SafeDownCast(sideset->GetPointData()->GetGlobalIds());

  if (pntIds == nullptr)
  {
    smtkErrorMacro(log, ent->name() + " global point ids are missing.");
    return false;
  }

  xmlNode.set_name("Surface");
  xmlNode.append_attribute("name") = ent->name().c_str();
  std::set<smtk::model::Entity*> primaries;
  session->primaryGeometries(ent, primaries);
  // track if a vertex cell is processed. By process, I mean its corresponding ids of its points are found in a primary volume mesh
  std::vector<bool> found(sideset->GetNumberOfCells(), false);
  // track how many vertex cells are processed.
  int numProcessed = 0;

  // since there may be multiple primaries for a single side-set, we find the corresponding ids of points primary-by-primary
  for (auto* surface : primaries)
  {
    if (!surface->isFace())
    {
      continue;
    }

    vtkSmartPointer<vtkPointSet> surface_mesh(
      vtkPointSet::SafeDownCast(session->findStorage(surface->id())));
    if (!surface_mesh)
    {
      smtkErrorMacro(log, surface->name() + " Could not convert to vtkPointSet.");
      continue;
    }
    vtkIntArray* surfacePntIds =
      vtkIntArray::SafeDownCast(surface_mesh->GetPointData()->GetGlobalIds());

    if (pntIds == nullptr)
    {
      smtkErrorMacro(log, "global point ids are missing.");
      return false;
    }

    int minGlobalId;
    int maxGlobalId;
    int minGlobalIdVol;
    int maxGlobalIdVol;
    getRangeOfGlobalIds<true>(session, surface->id(), minGlobalId, maxGlobalId);

    int febOffset;
    auto volIds = surface->relations();
    // If the volume geometry is empty or doesn't exist, replace it with the boundary
    // geometry and pretend there is a non-empty one.
    if (volIds.empty())
    {
      minGlobalIdVol = minGlobalId;
      maxGlobalIdVol = maxGlobalId;
      febOffset = pntOffset.at(surface->id());
    }
    else
    {
      auto foundVolRange =
        getRangeOfGlobalIds<true>(session, volIds[0], minGlobalIdVol, maxGlobalIdVol);
      if (!foundVolRange)
      {
        minGlobalIdVol = minGlobalId;
        maxGlobalIdVol = maxGlobalId;
        febOffset = pntOffset.at(surface->id());
      }
      else
      {
        febOffset = pntOffset.at(volIds[0]);
      }
    }

    // For a point of a cell in the side-set, we
    // 1. find its corresponding point id in its primary
    // 2. since this primary is extracted boundary of a volume, we have to find
    //    the corresponding id of this point in the volume
    // 3. convert the global id in the volume to feb id
    for (vtkIdType ii = 0; ii < sideset->GetNumberOfCells(); ++ii)
    {
      vtkNew<vtkIdList> pnts;
      // vtkIdList* pnts;
      sideset->GetCellPoints(ii, pnts);
      vtkIdType numPnts = pnts->GetNumberOfIds();
      std::string delimiter;

      std::string indices;
      int cnt = 0;
      for (int jj = 0; jj < numPnts; ++jj)
      {
        int id = pntIds->GetValue(pnts->GetId(jj));
        if (id < minGlobalId || id > maxGlobalId)
        {
          continue;
        }

        id += -minGlobalId;                        // local id in surface
        int idInVol = surfacePntIds->GetValue(id); // local id in surface to global id in volume
        if (idInVol < minGlobalIdVol || idInVol > maxGlobalIdVol)
        {
          smtkErrorMacro(log, "Cannot find surface point in the volume geometry.");
          continue;
        }
        idInVol += -minGlobalIdVol + febOffset + 1; // global id in volume to feb id

        found[ii] = true;
        ++cnt;
        delimiter = (jj == 0) ? " " : ", ";
        indices += delimiter + std::to_string(idInVol);
      }
      // all points should come from the same surface
      if (cnt < numPnts)
      {
        continue;
      }

      pugi::xml_node face = xmlNode.append_child(std::to_string(ii).c_str());
      face.set_name(cellNames[sideset->GetCellType(ii)].c_str());
      face.append_attribute("id") = std::to_string(ii + 1).c_str();
      face.append_child(pugi::node_pcdata).set_value(indices.c_str());
      ++numProcessed;
    }
  }

  if (numProcessed != sideset->GetNumberOfCells())
  {
    std::string msg = "side-set " + ent->name() + ": " + std::to_string(numProcessed) + " out of " +
      std::to_string(sideset->GetNumberOfCells()) + " found in primaries.";
    smtkErrorMacro(log, msg);
    return false;
  }

  return true;
}

bool addElementSet(smtk::model::Entity* ent,
  const smtk::session::aeva::Session::Ptr& session,
  smtk::io::Logger& log,
  const Offsets& cellOffset,
  pugi::xml_node& xmlNode)
{
  // convert to vtk mesh
  vtkSmartPointer<vtkPointSet> sideset(vtkPointSet::SafeDownCast(session->findStorage(ent->id())));
  if (!sideset)
  {
    smtkErrorMacro(log, "Could not convert to vtkPointSet.");
    return false;
  }
  vtkIntArray* cellIds = vtkIntArray::SafeDownCast(sideset->GetCellData()->GetGlobalIds());
  if (cellIds == nullptr)
  {
    smtkErrorMacro(log, "global cell ids are missing.");
    return false;
  }

  xmlNode.set_name("ElementSet");
  xmlNode.append_attribute("name") = ent->name().c_str();

  int numProcessed = 0;
  std::vector<bool> found(sideset->GetNumberOfCells(), false);
  std::set<smtk::model::Entity*> primaries;
  session->primaryGeometries(ent, primaries);

  for (auto* primary : primaries)
  {
    int minGlobalId;
    int maxGlobalId;
    getRangeOfGlobalIds<false>(session, primary->id(), minGlobalId, maxGlobalId);
    int febOffset = cellOffset.at(primary->id());

    for (vtkIdType ii = 0; ii < sideset->GetNumberOfCells(); ++ii)
    {
      int id = cellIds->GetValue(ii);
      if (found[ii] || id < minGlobalId || id > maxGlobalId)
      {
        continue;
      }
      id += febOffset - minGlobalId + 1;
      found[ii] = true;
      ++numProcessed;

      pugi::xml_node elem = xmlNode.append_child(std::to_string(id).c_str());
      elem.set_name("elem");
      elem.append_attribute("id") = std::to_string(id).c_str();
    }
  }

  if (numProcessed != sideset->GetNumberOfCells())
  {
    std::string msg = "side-set " + ent->name() + ": " + std::to_string(numProcessed) + " out of " +
      std::to_string(sideset->GetNumberOfCells()) + " found in primaries.";
    smtkErrorMacro(log, msg);
    return false;
  }

  return true;
}
} // anounymous namespace

ExportFeb::Result ExportFeb::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.empty())
  {
    filename += ".feb";
  }
  if (ext != ".feb")
  {
    smtkErrorMacro(this->log(), "Only file extension .feb is accepted.");
    return this->createResult(smtk::operation::Operation::Outcome::FAILED);
  }

  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);

  auto assocs = this->parameters()->associations();

  // put all geometries owned by input models together
  smtk::model::EntityRefs entRefs;
  for (const auto& assoc : *assocs)
  {
    auto assocResource = std::dynamic_pointer_cast<smtk::session::aeva::Resource>(assoc);
    if (assocResource)
    {
      smtk::model::Models models =
        resource->entitiesMatchingFlagsAs<smtk::model::Models>(smtk::model::MODEL_ENTITY, false);
      // add each model's cells.
      for (const auto& model : models)
      {
        const auto cells = model.cells();
        entRefs.insert(cells.begin(), cells.end());
      }
    }
    else
    {
      auto ent = std::dynamic_pointer_cast<smtk::model::Entity>(assoc);
      if (ent->isModel())
      {
        auto model = smtk::model::Model(
          std::static_pointer_cast<smtk::session::aeva::Resource>(ent->resource()), ent->id());
        const auto cells = model.cells();
        entRefs.insert(cells.begin(), cells.end());
      }
    }
  }

  pugi::xml_document root;
  auto febio = root.append_child("febio_spec");
  febio.append_attribute("version") = "2.5";
  auto geom = febio.append_child("Geometry");
  Offsets pntOffset;
  Offsets cellOffset;
  int pntCnt = 0;

  //
  // add Nodes
  //
  for (const auto& entRef : entRefs)
  {
    auto ent = entRef.entityRecord();
    if (Session::isSideSet(*ent))
    {
      // Note we don't have to worry about the boundary surface extracted from a volume since it's not added to entRefs
      continue;
    }
    // For some reason, session can be different for different resources. Re-fetch the session.
    session =
      std::dynamic_pointer_cast<smtk::session::aeva::Resource>(entRef.entityRecord()->resource())
        ->session();

    vtkSmartPointer<vtkPointSet> volMesh(
      vtkPointSet::SafeDownCast(session->findStorage(ent->id())));
    vtkPointSet* mesh = volMesh.Get();
    auto entId = ent->id();
    if (!mesh) // If the volume is empty, use its boundary mesh instead.
    {
      auto faceEnt = smtk::model::Volume(ent).faces()[0].entityRecord();
      entId = faceEnt->id();
      mesh = vtkPointSet::SafeDownCast(session->findStorage(faceEnt->id()));
    }
    if (!mesh)
    {
      smtkWarningMacro(this->log(), ent->name() + ": could not convert to vtkPointSet.");
      continue;
    }
    auto xmlNode = geom.append_child(ent->name().c_str());
    addNodes(*mesh, ent->name(), pntCnt, xmlNode);
    // record the point id offset of this primary volume mesh ofin .feb
    pntOffset[entId] = pntCnt;
    pntCnt += mesh->GetNumberOfPoints();
  }

  //
  // add Elements
  //
  int cellCnt = 0;
  for (const auto& entRef : entRefs)
  {
    auto ent = entRef.entityRecord();
    if (Session::isSideSet(*ent))
    {
      continue;
    }

    vtkSmartPointer<vtkPointSet> volMesh(
      vtkPointSet::SafeDownCast(session->findStorage(ent->id())));
    vtkPointSet* mesh = volMesh.Get();
    auto entId = ent->id();
    if (!mesh) // If the volume is empty, use its boundary mesh instead.
    {
      // auto& faceEnt = dynamic_cast<const smtk::model::Volume&>(entRef).faces()[0];
      auto faceEnt = smtk::model::Volume(ent).faces()[0].entityRecord();
      mesh = vtkPointSet::SafeDownCast(session->findStorage(faceEnt->id()));
      entId = faceEnt->id();
    }
    if (!mesh)
    {
      smtkWarningMacro(this->log(), ent->name() + ": could not convert to vtkPointSet.");
      continue;
    }

    auto xmlNode = geom.append_child(ent->name().c_str());
    addElements(*mesh, ent->name(), pntOffset.at(entId), cellCnt, xmlNode);
    cellOffset[entId] = cellCnt;
    cellCnt += mesh->GetNumberOfCells();
  }

  // Add node sets
  for (const auto& entRef : entRefs)
  {
    auto ent = entRef.entityRecord();
    if (Session::isPrimary(*ent) || !ent->isVertex())
    {
      continue;
    }

    auto xmlNode = geom.append_child(ent->name().c_str());
    addNodeSet(ent.get(), session, this->log(), pntOffset, xmlNode);
  }

  // Add face sets
  for (const auto& entRef : entRefs)
  {
    auto ent = entRef.entityRecord();
    if (Session::isPrimary(*ent) || !ent->isFace())
    {
      continue;
    }

    auto xmlNode = geom.append_child(ent->name().c_str());
    addFaceSet(ent.get(), session, this->log(), pntOffset, xmlNode);
  }

  // Add element sets
  for (const auto& entRef : entRefs)
  {
    auto ent = entRef.entityRecord();
    if (Session::isPrimary(*ent) || !ent->isVolume())
    {
      continue;
    }

    auto xmlNode = geom.append_child(ent->name().c_str());
    addElementSet(ent.get(), session, this->log(), cellOffset, xmlNode);
  }

  root.save_file(filename.c_str());
  return this->createResult(smtk::operation::Operation::Outcome::SUCCEEDED);
}

const char* ExportFeb::xmlDescription() const
{
  return ExportFeb_xml;
}

}
}
}
