//=============================================================================
// 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/CellSelection.h"
#include "smtk/session/aeva/Predicates.h"
#include "smtk/session/aeva/operators/GrowSelection.h"

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

#include "smtk/attribute/Attribute.h"
#include "smtk/attribute/ComponentItem.h"
#include "smtk/attribute/DoubleItem.h"
#include "smtk/attribute/IntItem.h"
#include "smtk/attribute/VoidItem.h"

#include "smtk/operation/MarkGeometry.h"

#include "vtkAppendDataSets.h"
#include "vtkCellData.h"
#include "vtkDataSetSurfaceFilter.h"
#include "vtkDistancePolyDataFilter.h"
#include "vtkDoubleArray.h"
#include "vtkImplicitPolyDataDistance.h"
#include "vtkMath.h"
#include "vtkPointData.h"
#include "vtkPolyData.h"
#include "vtkPolyDataNormals.h"
#include "vtkThreshold.h"
#include "vtkUnsignedCharArray.h"
#include "vtkUnstructuredGrid.h"
#include "vtkVector.h"
#include "vtkVectorOperators.h"

#include <cmath>

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

bool GrowSelection::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;
  }

  auto assocs = this->parameters()->associations();
  auto seed = std::dynamic_pointer_cast<smtk::resource::Component>(assocs->value());
  return seed && !session->isPrimary(*seed) && !!Operation::storage(seed);
}

bool GrowSelection::growFromSideSetSeed(const smtk::model::Entity::Ptr& seed,
  double angleTol,
  Operation::Result& result,
  smtk::io::Logger& log)
{
  if (!seed)
  {
    smtkErrorMacro(log, "Source must be a primitive selection or side set with geometry.");
    return false;
  }
  auto resource = dynamic_pointer_cast<smtk::session::aeva::Resource>(seed->resource());
  Session::Ptr session = resource ? resource->session() : Session::Ptr();
  if (!session)
  {
    smtkErrorMacro(log, "Source must have a valid resource and session.");
    return false;
  }
  vtkSmartPointer<vtkDataObject> seedData;
  if (Session::isPrimary(*seed) || !(seedData = Operation::storage(seed)))
  {
    smtkErrorMacro(log, "Source must be a primitive selection or side set with geometry.");
    return false;
  }

  // I. Find parent primary geometry that CellSelection is drawn from.
  //    May be multiple cells (e.g., a volume bounded by several surfaces).
  std::set<smtk::model::Entity*> primaries;
  if (!session->primaryGeometries(seed.get(), primaries))
  {
    smtkErrorMacro(log, "Could not find primary geometry for input " << seed->name() << ".");
    return false;
  }
  // II. Find matching cells in primary geometry.
  std::map<smtk::model::Entity*, std::set<vtkIdType> > seedMap;
  int entityType = -1;
  if (!GrowSelection::findSeedsInPrimaries(seedData, primaries, seedMap, entityType, log))
  {
    smtkErrorMacro(log, "Could not find seeds in primary geometry.");
    return false;
  }
  if (entityType != vtkDataObject::CELL)
  {
    smtkErrorMacro(log, "Selected cells were points, Cannot grow points by angle.");
    return false;
  }
  // III. Traverse cells to grow, appending to CellSelection as needed.
  vtkSmartPointer<vtkDataObject> updated =
    GrowSelection::growFromSeeds(seedData, seedMap, angleTol);
  if (updated == seedData)
  {
    smtkWarningMacro(log, "No cells added to selection.");
  }
  else
  {
    session->addStorage(seed->id(), updated);
    smtk::operation::MarkGeometry().markModified(seed);
    result->findComponent("modified")->appendValue(seed);
  }
  return true;
}

bool GrowSelection::findSeedsInPrimaries(const vtkSmartPointer<vtkDataObject>& seedData,
  std::set<smtk::model::Entity*>& primaries,
  std::map<smtk::model::Entity*, std::set<vtkIdType> >& seedMap,
  int& entityType,
  smtk::io::Logger& log)
{
  if (!seedData)
  {
    return false;
  }
  entityType = static_cast<int>(EntityType(seedData));
  auto* attr = seedData->GetAttributes(entityType);
  if (!attr)
  {
    smtkErrorMacro(log, "No attributes of type " << entityType << ".");
    return false;
  }
  auto* gids = vtkIntArray::SafeDownCast(attr->GetGlobalIds());
  if (!gids)
  {
    smtkErrorMacro(log, "No global IDs.");
    return false;
  }
  if (primaries.empty())
  {
    smtkErrorMacro(log, "No matching primary geometry.");
    return false;
  }
  vtkIdType ng = gids->GetNumberOfValues();
  std::set<vtkIdType> ids;
  for (vtkIdType ii = 0; ii < ng; ++ii)
  {
    ids.insert(gids->GetValue(ii));
  }
  // Now iterate over primary geometry and record matching points/cells.
  bool didFind = false;
  for (auto* primary : primaries)
  {
    auto pdata = Operation::storage(primary);
    if (!pdata)
    {
      smtkWarningMacro(log, "Primary \"" << primary->name() << "\" has no geometry.");
      continue;
    }
    auto* patt = pdata->GetAttributes(entityType);
    if (!patt)
    {
      continue;
    }
    auto* pgids = vtkIntArray::SafeDownCast(patt->GetGlobalIds());
    if (!pgids)
    {
      continue;
    }
    auto smIt = seedMap.find(primary);
    if (smIt == seedMap.end())
    {
      smIt =
        seedMap.insert(std::pair<smtk::model::Entity*, std::set<vtkIdType> >(primary, {})).first;
    }
    std::set<vtkIdType>& entries(smIt->second);
    vtkIdType nprim = pgids->GetNumberOfValues();
    for (vtkIdType ii = 0; ii < nprim; ++ii)
    {
      if (ids.find(pgids->GetValue(ii)) != ids.end())
      {
        entries.insert(ii);
        didFind = true;
      }
    }
  }
  // TODO: Could verify that ids.size() == number of times didFind set true above.
  return didFind;
}

vtkSmartPointer<vtkDataObject> GrowSelection::growOnPrimary(
  const vtkSmartPointer<vtkDataSet>& primary,
  const std::set<vtkIdType>& seeds,
  double angleTol)
{
  vtkSmartPointer<vtkPolyData> surf;
  vtkSmartPointer<vtkDataArray> normals;
  SurfaceWithNormals(primary, surf, normals);

  auto edgeVisitor = EdgeNeighbors(surf);
  std::set<vtkIdType> queue{ seeds };
  // Create an array to hold markup on cells:
  //   0 - unvisited
  //   1 - visited but not selected
  //   2 - visited and chosen for selection
  vtkNew<vtkUnsignedCharArray> visited;
  visited->SetName("SelectionValue");
  visited->SetNumberOfTuples(normals->GetNumberOfTuples());
  visited->FillComponent(0, 0);

  while (!queue.empty())
  {
    auto it = queue.begin();
    vtkIdType cellId = *it;
    if (visited->GetValue(cellId) < 2)
    {
      // Get seed cell normal
      vtkVector3d cellNormal;
      normals->GetTuple(cellId, cellNormal.GetData());
      visited->SetValue(cellId, 2);
      edgeVisitor(cellId, [angleTol, &visited, &queue, &cellNormal, &normals](vtkIdType neighbor) {
        if (visited->GetValue(neighbor) < 2)
        {
          vtkVector3d neighborNormal;
          normals->GetTuple(neighbor, neighborNormal.GetData());
          double dot = cellNormal.Dot(neighborNormal);
          if (dot > angleTol)
          {
            queue.insert(neighbor);
          }
          visited->SetValue(neighbor, 1);
        }
        return Continue;
      });
    }
    visited->SetValue(cellId, 2);
    queue.erase(it);
  }
  // Now turn visited into a side set.
  vtkNew<vtkThreshold> threshold;
  surf->GetCellData()->AddArray(visited);
  threshold->SetInputDataObject(surf);
  threshold->SetInputArrayToProcess(
    0, 0, 0, vtkDataObject::FIELD_ASSOCIATION_CELLS, "SelectionValue");
  threshold->ThresholdBetween(2, 255);
  threshold->Update();

  auto* grown = threshold->GetOutput();
  grown->GetCellData()->RemoveArray("SelectionValue");
  grown->GetCellData()->RemoveArray("Normals");
  grown->GetPointData()->RemoveArray("Normals");
  return grown;
}

vtkSmartPointer<vtkDataObject> GrowSelection::growFromSeeds(vtkSmartPointer<vtkDataObject> seedData,
  std::map<smtk::model::Entity*, std::set<vtkIdType> >& seedMap,
  double angleTol)
{
  vtkSmartPointer<vtkDataObject> updated;
  vtkNew<vtkAppendDataSets> append;
  for (const auto& entry : seedMap)
  {
    auto* primary = vtkDataSet::SafeDownCast(Operation::storage(entry.first));
    if (primary)
    {
      auto grown = GrowSelection::growOnPrimary(primary, entry.second, angleTol);
      append->AddInputData(grown);
    }
  }
  append->Update();
  updated = append->GetOutputDataObject(0);

  // If we haven't grown, indicate it by returning the input.
  // This prevents us from re-running the visualization pipeline.
  if (NumberOfCells(updated) <= NumberOfCells(seedData))
  {
    return seedData;
  }
  return updated;
}

GrowSelection::Result GrowSelection::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);

  auto angleItem = this->parameters()->findDouble("angle");
  double angle = angleItem->value();
  double angleTol = std::cos(vtkMath::RadiansFromDegrees(angle));

  auto assocs = this->parameters()->associations();
  auto seed = std::dynamic_pointer_cast<smtk::model::Entity>(assocs->value());
  if (GrowSelection::growFromSideSetSeed(seed, angleTol, result, this->log()))
  {
    result->findInt("outcome")->setValue(static_cast<int>(GrowSelection::Outcome::SUCCEEDED));
  }
  return result;
}

const char* GrowSelection::xmlDescription() const
{
  return GrowSelection_xml;
}

}
}
}
