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

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

#include "smtk/model/Entity.h"

#include "vtkCellData.h"
#include "vtkCellType.h"
#include "vtkDataObject.h"
#include "vtkDataSet.h"
#include "vtkGenerateGlobalIds.h"
#include "vtkIntArray.h"
#include "vtkPointData.h"

namespace smtk
{
namespace session
{
namespace aeva
{

constexpr unsigned int Session::Dimension;
constexpr const char* const Session::type_name;

std::string Session::defaultFileExtension(const smtk::model::Model& /*model*/) const
{
  return ".vti";
}

vtkSmartPointer<vtkDataObject> Session::findStorage(const smtk::common::UUID& uid) const
{
  auto it = m_storage.find(uid);
  if (it != m_storage.end())
  {
    return it->second;
  }
  return nullptr;
}

void Session::addStorage(const smtk::common::UUID& uid,
  vtkSmartPointer<vtkDataObject> const& storage)
{
  if (uid.isNull())
  {
    return;
  }

  if (!storage)
  {
    m_storage.erase(uid);
  }
  else
  {
    m_storage[uid] = storage;
  }
}

#if 0
Session::ImageType::Pointer Session::findStorage(const smtk::common::UUID& uid) const
{
  auto it = m_storage.find(uid);
  if (it != m_storage.end())
  {
    return it->second;
  }
  return nullptr;
}

void Session::addStorage(const smtk::common::UUID& uid, ImageType::Pointer storage)
{
  if (uid.isNull())
  {
    return;
  }

  if (!storage)
  {
    m_storage.erase(uid);
  }
  else
  {
    m_storage[uid] = storage;
  }
}
#endif

bool Session::removeStorage(const smtk::common::UUID& uid)
{
  if (uid.isNull())
  {
    return false;
  }

  return m_storage.erase(uid) > 0;
}

bool Session::primaryGeometries(const vtkSmartPointer<vtkDataObject>& data,
  int fieldType, // one of vtkDataObject::AttributeTypes
  std::set<smtk::model::Entity*>& primaries,
  smtk::resource::Resource* resource) const
{
  vtkDataSetAttributes* attr;
  vtkIntArray* gids;
  if (!resource || !data || !(attr = data->GetAttributes(fieldType)) ||
    !(gids = vtkIntArray::SafeDownCast(attr->GetGlobalIds())))
  {
    return false;
  }

  vtkIdType nn = gids->GetNumberOfValues();
  int* tmp = gids->GetValueRange(0);
  std::array<int, 2> dataIdRange{ tmp[0], tmp[1] };

  bool didAdd = false;
  for (const auto& it : m_storage)
  {
    auto comp = std::dynamic_pointer_cast<smtk::model::Entity>(resource->find(it.first));
    if (!comp || !Session::isPrimary(*comp))
    {
      continue;
    }
    // FIXME.
    // Primary objects have contiguous global cell ID ranges
    // but not necessarily point IDs (example: primary surface
    // and volume meshes may share points). Maybe create a
    // "vertex" cell marked primary guaranteed to have contiguous
    // point IDs, which we should always prefer when fieldType is
    // POINTS?
    auto* primaryIds =
      vtkIntArray::SafeDownCast(it.second->GetAttributes(fieldType)->GetGlobalIds());
    if (!primaryIds)
    {
      continue;
    }
    tmp = primaryIds->GetValueRange(0);
    std::array<int, 2> primIdRange{ tmp[0], tmp[1] };
    if (primIdRange[1] < dataIdRange[0] || primIdRange[0] > dataIdRange[1])
    {
      continue;
    }
    // ID ranges overlap. We must handle a special case.
    if (primIdRange[0] < dataIdRange[0] && primIdRange[1] > dataIdRange[1])
    {
      // VTK data may have IDs from primary data "before" and "after"
      // the current primary data, but not the current object.
      // Verify by examining every ID until we find at least one in range.
      bool inRange = false;
      for (vtkIdType ii = 0; ii < nn; ++ii)
      {
        int id = gids->GetValue(ii);
        if (id >= primIdRange[0] && id <= primIdRange[1])
        {
          inRange = true;
          break;
        }
      }
      if (!inRange)
      {
        continue;
      }
    }
    didAdd = true;
    primaries.insert(comp.get());
  }

  return didAdd;
}

bool Session::primaryGeometries(const smtk::model::Entity* in,
  std::set<smtk::model::Entity*>& primaries) const
{
  if (!in)
  {
    return false;
  }
  auto geom = this->findStorage(in->id());
  if (!geom)
  {
    return false;
  }
  int fieldType = in->dimension() == 0 ? vtkDataObject::POINT : vtkDataObject::CELL;
  return this->primaryGeometries(geom, fieldType, primaries, in->resource().get());
}

smtk::model::Entity* Session::primaryGeometry(const smtk::model::Entity* in) const
{
  std::set<smtk::model::Entity*> primaries;
  if (!this->primaryGeometries(in, primaries) || primaries.size() != 1)
  {
    return nullptr;
  }
  return *primaries.begin();
}

smtk::common::UUID Session::primaryGeometryForCell(vtkIdType globalId) const
{
  for (const auto& entry : m_storage)
  {
    auto* attr = entry.second->GetAttributes(vtkDataObject::CELL);
    if (attr)
    {
      auto* gids = vtkIntArray::SafeDownCast(attr->GetGlobalIds());
      if (gids)
      {
        int* range = gids->GetValueRange();
        if (range[0] <= globalId && globalId <= range[1])
        {
          return entry.first;
        }
      }
    }
  }
  return smtk::common::UUID::null();
}

smtk::common::UUID Session::primaryGeometryForPoint(vtkIdType globalId) const
{
  for (const auto& entry : m_storage)
  {
    auto* attr = entry.second->GetAttributes(vtkDataObject::POINT);
    if (attr)
    {
      auto* gids = vtkIntArray::SafeDownCast(attr->GetGlobalIds());
      if (gids)
      {
        int* range = gids->GetValueRange();
        if (range[0] <= globalId && globalId <= range[1])
        {
          return entry.first;
        }
      }
    }
  }
  return smtk::common::UUID::null();
}

std::map<vtkIdType, Session::DataEntry> Session::buildGlobalIdLookup(
  const std::shared_ptr<Resource>& resource,
  vtkDataObject::AttributeTypes centering) const
{
  std::map<vtkIdType, DataEntry> result;
  for (const auto& entry : m_storage)
  {
    auto comp = resource->find(entry.first);
    if (!comp)
    {
      continue;
    }
    if (!Session::isPrimary(*comp))
    {
      continue;
    }
    auto* attr = entry.second->GetAttributes(centering);
    if (attr)
    {
      auto* gids = vtkIntArray::SafeDownCast(attr->GetGlobalIds());
      if (gids)
      {
        int* range = gids->GetValueRange();
        result[range[0]] = entry;
        result[range[1]] = entry;
      }
    }
  }
  return result;
}

bool Session::sideSets(const smtk::model::Entity* in,
  std::set<smtk::model::Entity*>& sideSets) const
{
  bool didAdd = false;
  if (!in)
  {
    return didAdd;
  }

  auto resource = std::dynamic_pointer_cast<Resource>(in->resource());
  if (!resource)
  {
    return didAdd;
  }

  for (const auto& it : m_storage)
  {
    auto comp = std::dynamic_pointer_cast<smtk::model::Entity>(resource->find(it.first));
    if (!comp || Session::isPrimary(*comp))
    {
      continue;
    }

    int sideSetDim = EstimateParametricDimension(it.second);
    auto fieldType = sideSetDim > 0 ? vtkDataObject::CELL : vtkDataObject::POINT;
    std::set<smtk::model::Entity*> primaries;
    if (this->primaryGeometries(it.second, fieldType, primaries, resource.get()))
    {
      if (primaries.find(const_cast<smtk::model::Entity*>(in)) != primaries.end())
      {
        sideSets.insert(comp.get());
        didAdd = true;
      }
    }
  }

  return didAdd;
}

smtk::model::SessionInfoBits Session::transcribeInternal(const smtk::model::EntityRef& /*entity*/,
  SessionInfoBits /*requestedInfo*/,
  int /*depth*/)
{
  return smtk::model::SESSION_EVERYTHING;
}

smtk::model::SessionIOPtr Session::createIODelegate(const std::string& /*format*/)
{
  // TODO: If we write any custom JSON to .smtk files, return a handler.
  return nullptr;
}

vtkVisitation Session::visitPrimaryDataById(const std::shared_ptr<Resource>& resource,
  const std::set<vtkIdType>& ids,
  vtkVisitationRule context,
  const ElementVisitor& visitor) const
{
  bool restrictive = false;
  switch (context)
  {
    case vtkVisitationRule::CELLS_COMPOSED_OF_POINTS:
      restrictive = true;
      // fall through
    case vtkVisitationRule::CELLS_CONNECTED_TO_POINTS:
    {
      vtkNew<vtkIdList> cellsUsingPoint;
      vtkNew<vtkIdList> cellConnectivity;
      std::set<int> visited;
      auto index = this->buildGlobalIdLookup(resource, vtkDataObject::POINT);
      for (const auto& pid : ids)
      {
        const auto& indexIt = index.lower_bound(pid);
        if (indexIt == index.end())
        {
          continue;
        }
        auto* data = vtkDataSet::SafeDownCast(indexIt->second.second);
        if (!data)
        {
          continue;
        }
        auto* dataPointIds =
          vtkIntArray::SafeDownCast(data->GetAttributes(vtkDataObject::POINT)->GetGlobalIds());
        if (!dataPointIds)
        {
          continue;
        }
        auto* dataCellIds =
          vtkIntArray::SafeDownCast(data->GetAttributes(vtkDataObject::CELL)->GetGlobalIds());
        if (!dataCellIds)
        {
          continue;
        }
        int* range = dataPointIds->GetValueRange();
        int offset = static_cast<int>(pid - static_cast<vtkIdType>(range[0]));
        if (offset < 0 || offset >= data->GetNumberOfPoints())
        {
          continue;
        }
        data->GetPointCells(offset, cellsUsingPoint);
        for (const auto& cellOffset : *cellsUsingPoint)
        {
          int cellId = dataCellIds->GetValue(cellOffset);
          if (visited.find(cellId) != visited.end())
          {
            continue;
          }
          visited.insert(cellId);
          data->GetCellPoints(cellOffset, cellConnectivity);
          if (restrictive)
          {
            // Verify all points have global IDs in set:
            bool passed = true;
            for (const auto& connEntry : *cellConnectivity)
            {
              int testGlobalId = dataPointIds->GetValue(connEntry);
              if (ids.find(testGlobalId) == ids.end())
              {
                passed = false;
                break;
              }
            }
            if (!passed)
            {
              continue;
            }
          }
          int cellType = data->GetCellType(cellOffset);
          vtkVisitation proceed = visitor(cellId,
            pid,
            data,
            cellType,
            cellConnectivity->GetNumberOfIds(),
            cellConnectivity->GetPointer(0));
          if (proceed == vtkVisitation::STOP)
          {
            return proceed;
          }
        }
      }
      return vtkVisitation::CONTINUE;
    }
    break;
    case vtkVisitationRule::POINTS_INTERNAL_TO_CELLS:
      restrictive = true;
      // fall through
    case vtkVisitationRule::POINTS_CONNECTED_TO_CELLS:
    {
      std::cerr << "ERROR: Not implemented.\n";
      vtkNew<vtkIdList> cellsUsingPoint;
      vtkNew<vtkIdList> cellConnectivity;
      std::set<int> visited;
      auto index = this->buildGlobalIdLookup(resource, vtkDataObject::CELL);
      for (const auto& cid : ids)
      {
        const auto& indexIt = index.lower_bound(cid);
        if (indexIt == index.end())
        {
          continue;
        }
        auto* data = vtkDataSet::SafeDownCast(indexIt->second.second);
        if (!data)
        {
          continue;
        }
        auto* dataPointIds =
          vtkIntArray::SafeDownCast(data->GetAttributes(vtkDataObject::POINT)->GetGlobalIds());
        if (!dataPointIds)
        {
          continue;
        }
        auto* dataCellIds =
          vtkIntArray::SafeDownCast(data->GetAttributes(vtkDataObject::CELL)->GetGlobalIds());
        if (!dataCellIds)
        {
          continue;
        }
        int* range = dataCellIds->GetValueRange();
        int offset = static_cast<int>(cid - static_cast<vtkIdType>(range[0]));
        if (offset < 0 || offset >= data->GetNumberOfCells())
        {
          continue;
        }
        data->GetCellPoints(offset, cellConnectivity);
        for (const auto& pointOffset : *cellConnectivity)
        {
          int pointId = dataPointIds->GetValue(pointOffset);
          if (visited.find(pointId) != visited.end())
          {
            continue;
          }
          visited.insert(pointId);
          data->GetPointCells(pointOffset, cellsUsingPoint);
          vtkIdType someCell =
            cellsUsingPoint->GetNumberOfIds() > 0 ? cellsUsingPoint->GetId(0) : -1;
          if (restrictive)
          {
            // Verify all points have global IDs in set:
            bool passed = true;
            for (const auto& cellAtCorner : *cellsUsingPoint)
            {
              int testGlobalId = dataCellIds->GetValue(cellAtCorner);
              if (ids.find(testGlobalId) == ids.end())
              {
                passed = false;
                break;
              }
            }
            if (!passed)
            {
              continue;
            }
          }
          int cellType = VTK_VERTEX;
          vtkVisitation proceed = visitor(pointId, someCell, data, cellType, 1, &pointOffset);
          if (proceed == vtkVisitation::STOP)
          {
            return proceed;
          }
        }
      }
      return vtkVisitation::CONTINUE;
    }
    break;
  }
  return vtkVisitation::STOP;
}

void Session::offsetGlobalIds(vtkDataSet* data,
  long pointIdOffset,
  long cellIdOffset,
  long& maxPointId,
  long& maxCellId)
{
  if (!data)
  {
    return;
  }
  vtkIntArray* pointDataNums = vtkIntArray::SafeDownCast(data->GetPointData()->GetGlobalIds());
  vtkIntArray* cellDataNums = vtkIntArray::SafeDownCast(data->GetCellData()->GetGlobalIds());

  if (!pointDataNums || !cellDataNums)
  {
    // need to generate globalIds.
    vtkNew<vtkGenerateGlobalIds> generator;
    generator->SetInputDataObject(data);
    generator->Update();
    auto* outdata = vtkDataSet::SafeDownCast(generator->GetOutputDataObject(0));

    // transform point to ints, to matcch .med format
    if (!pointDataNums)
    {
      vtkNew<vtkIntArray> newPointIds;
      newPointIds->DeepCopy(outdata->GetPointData()->GetGlobalIds());
      data->GetPointData()->SetGlobalIds(newPointIds);
      pointDataNums = vtkIntArray::SafeDownCast(data->GetPointData()->GetGlobalIds());
    }

    if (!cellDataNums)
    {
      vtkNew<vtkIntArray> newCellIds;
      newCellIds->DeepCopy(outdata->GetCellData()->GetGlobalIds());
      data->GetCellData()->SetGlobalIds(newCellIds);
      cellDataNums = vtkIntArray::SafeDownCast(data->GetCellData()->GetGlobalIds());
    }
  }

  if (pointDataNums && pointDataNums->GetNumberOfValues() > 0)
  {
    bool doOffset = pointDataNums->GetValue(0) <= pointIdOffset;
    for (int i = 0; i < pointDataNums->GetNumberOfValues(); ++i)
    {
      if (doOffset)
      {
        pointDataNums->SetValue(i, pointDataNums->GetValue(i) + pointIdOffset);
      }
      maxPointId = std::max(maxPointId, static_cast<long>(pointDataNums->GetValue(i)));
    }
  }

  if (cellDataNums && cellDataNums->GetNumberOfValues() > 0)
  {
    bool doOffset = cellDataNums->GetValue(0) <= cellIdOffset;
    for (int i = 0; i < cellDataNums->GetNumberOfValues(); ++i)
    {
      if (doOffset)
      {
        cellDataNums->SetValue(i, cellDataNums->GetValue(i) + cellIdOffset);
      }
      maxCellId = std::max(maxCellId, static_cast<long>(cellDataNums->GetValue(i)));
    }
  }
}

void Session::setPrimary(smtk::resource::Component& c, bool isPrimary)
{
  if (isPrimary)
  {
    c.properties().get<long>()["primary"] = 1;
    c.properties().erase<long>("side-set");
  }
  else
  {
    c.properties().get<long>()["side-set"] = 1;
    c.properties().erase<long>("primary");
  }
}

bool Session::isPrimary(const smtk::resource::Component& c)
{
  if (c.properties().contains<long>("primary"))
  {
    return !!c.properties().at<long>("primary");
  }
  return false;
}

vtkIntArray* Session::globalIds(vtkDataObject* data, vtkDataObject::AttributeTypes centering)
{
  if (!data)
  {
    return nullptr;
  }
  auto* dsa = data->GetAttributes(centering);
  if (!dsa)
  {
    return nullptr;
  }
  // aeva's global IDs are currently forced to be integers.
  // Do not accept IDs which are not.
  auto* igid = vtkIntArray::SafeDownCast(dsa->GetGlobalIds());
  if (!igid)
  {
    for (int ii = 0; ii < dsa->GetNumberOfArrays(); ++ii)
    {
      igid = vtkIntArray::SafeDownCast(dsa->GetArray(ii));
      if (igid)
      {
        break;
      }
    }
  }
  return igid;
}

} // namespace aeva
} // namespace session
} // namespace smtk
