//=========================================================================
//  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/rgg/plugin/rggAuxiliaryGeometryExtension.h"
#include "smtk/extension/vtk/source/vtkCmbLayeredConeSource.h"
#include "smtk/extension/vtk/model/vtkModelAuxiliaryGeometry.h"
#include "smtk/extension/vtk/model/vtkModelAuxiliaryGeometry.txx"
#include "smtk/extension/vtk/source/vtkModelMultiBlockSource.h"

#include "smtk/AutoInit.h"
#include "smtk/PublicPointerDefs.h"

#include "smtk/model/AuxiliaryGeometry.h"
#include "smtk/model/Model.h"
#include "smtk/model/Resource.h"

#include "smtk/session/rgg/Core.h"
#include "smtk/session/rgg/Duct.h"
#include "smtk/session/rgg/Pin.h"
#include "smtk/session/rgg/operators/OperationHelper.h"

#include "smtk/session/rgg/json/jsonCore.h"
#include "smtk/session/rgg/json/jsonDuct.h"
#include "smtk/session/rgg/json/jsonPin.h"

#include "vtkBoundingBox.h"
#include "vtkClipClosedSurface.h"
#include "vtkCompositeDataSet.h"
#include "vtkDataObject.h"
#include "vtkDataObjectTreeIterator.h"
#include "vtkDataSet.h"
#include "vtkDoubleArray.h"
#include "vtkFloatArray.h"
#include "vtkGraph.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkNew.h"
#include "vtkPlaneCollection.h"
#include "vtkPointData.h"
#include "vtkPoints.h"
#include "vtkPolyData.h"
#include "vtkPolyDataNormals.h"
#include "vtkProperty.h"
#include "vtkStringArray.h"
#include "vtkTransform.h"
#include "vtkTransformPolyDataFilter.h"
#include "vtkUnstructuredGrid.h"

#include <ctime>
#include <list>
#include <map>

using namespace smtk::model;
using namespace smtk::session::rgg;
using rggPin = smtk::session::rgg::Pin;
using rggDuct = smtk::session::rgg::Duct;
using json = nlohmann::json;

namespace
{
const double cos30 = 0.86602540378443864676372317075294;
static void normalize(double* xyz)
{
  double sum = std::sqrt(xyz[0] * xyz[0] + xyz[1] * xyz[1] + xyz[2] * xyz[2]);
  if (sum == 0.0)
    return;
  xyz[0] /= sum;
  xyz[1] /= sum;
  xyz[2] /= sum;
}

static void transformNormal(double* xyz, double* xformR)
{
  if (xformR[0] != 0.0 || xformR[1] != 0.0 || xformR[2] != 0.0)
  {
    vtkSmartPointer<vtkTransform> xform = vtkSmartPointer<vtkTransform>::New();
    xform->RotateX(xformR[0]);
    xform->RotateY(xformR[1]);
    xform->RotateZ(xformR[2]);
    double tp[3];
    xform->TransformNormal(xyz, tp);
    xyz[0] = tp[0];
    xyz[1] = tp[1];
    xyz[2] = tp[2];
    normalize(xyz);
  }
}

static void clip(vtkSmartPointer<vtkPolyData> in, vtkSmartPointer<vtkPolyData>& out, double* normal,
  int offset = 0)
{
  vtkSmartPointer<vtkPolyData> tmpIn = in;
  out = vtkSmartPointer<vtkPolyData>::New();
  vtkNew<vtkPlane> plane;
  plane->SetOrigin(
    -normal[0] * 0.005 * offset, -normal[1] * 0.005 * offset, -normal[2] * 0.005 * offset);
  plane->SetNormal(normal[0], normal[1], normal[2]);

  vtkNew<vtkClipClosedSurface> clipper;
  vtkNew<vtkPlaneCollection> clipPlanes;
  vtkNew<vtkPolyDataNormals> normals;
  clipPlanes->AddItem(plane.GetPointer());
  clipper->SetClippingPlanes(clipPlanes.GetPointer());
  clipper->SetActivePlaneId(0);
  clipper->SetClipColor(1.0, 1.0, 1.0);
  clipper->SetActivePlaneColor(1.0, 1.0, 0.8);
  clipper->GenerateOutlineOff();
  clipper->SetInputData(tmpIn);
  clipper->GenerateFacesOn();
  normals->SetInputConnection(clipper->GetOutputPort());
  normals->Update();
  out->DeepCopy(normals->GetOutput());
}
}

rggAuxiliaryGeometryExtension::rggAuxiliaryGeometryExtension()
{
  rggAuxiliaryGeometryExtension::ensureCache();
}

rggAuxiliaryGeometryExtension::~rggAuxiliaryGeometryExtension()
{
}

bool rggAuxiliaryGeometryExtension::canHandleAuxiliaryGeometry(
  smtk::model::AuxiliaryGeometry& entity, std::vector<double>& bboxOut)
{
  if (!entity.isValid() || !entity.hasStringProperty("rggType"))
  {
    return false;
  }

  auto dataset = vtkAuxiliaryGeometryExtension::fetchCachedGeometry(entity);

  // Since resource is using a map to store uuids, the order of whether pin or its sub auxgeoms
  // will be processed is undefined
  smtk::model::AuxiliaryGeometry toplevelAuxgeom = entity;

  if (dataset)
  {
    return this->updateBoundsFromDataSet(entity, bboxOut, dataset);
  }
  else if (!dataset && entity.embeddedIn().isAuxiliaryGeometry())
  {
    // We are fetching the child's geometry first. Let's ask its parent to generate the geometry
    toplevelAuxgeom = entity.embeddedIn();
  }

  // We might add multiple auxgeoms. Don't trim it.
  bool trimCache = false;

  Model parentModel = entity.owningModel();
  bool genNormals = false;
  if (parentModel.hasIntegerProperty("generate normals"))
  {
    const IntegerList& prop(parentModel.integerProperty("generate normals"));
    if (!prop.empty() && prop[0])
    {
      genNormals = true;
    }
  }

  dataset = rggAuxiliaryGeometryExtension::generateRGGRepresentation(toplevelAuxgeom, genNormals);
  std::time_t mtime;
  std::time(&mtime);
  this->addCacheGeometry(dataset, toplevelAuxgeom, mtime, trimCache);

  if (toplevelAuxgeom != entity)
  { // Fetch the geometry for the child auxgeom
    dataset = vtkAuxiliaryGeometryExtension::fetchCachedGeometry(entity);
  }
  return this->updateBoundsFromDataSet(entity, bboxOut, dataset);
}

vtkSmartPointer<vtkDataObject> rggAuxiliaryGeometryExtension::generateRGGRepresentation(
  const AuxiliaryGeometry& rggEntity, bool genNormals)
{
  rggAuxiliaryGeometryExtension::ensureCache();
  if (rggEntity.stringProperty("rggType")[0] == rggPin::typeDescription)
  {
    if (rggEntity.auxiliaryGeometries().size() == 0)
    {
      smtkErrorMacro(smtk::io::Logger::instance(), "Cannot create a representation for pin "
          << rggEntity.name() << "without any parts or layers. If itself is a part or layer,"
                                 "its parent should generate the rep for it");
      return vtkSmartPointer<vtkDataObject>();
    }
    return rggAuxiliaryGeometryExtension::generateRGGPinRepresentation(rggEntity, genNormals);
  }
  else if (rggEntity.stringProperty("rggType")[0] == smtk::session::rgg::Duct::typeDescription)
  {
    if (rggEntity.auxiliaryGeometries().size() == 0)
    {
      smtkErrorMacro(smtk::io::Logger::instance(), "Cannot create a representation for duct "
          << rggEntity.name() << "without any segments or layers. If itself is a segment or layer,"
                                 "its parent should generate the rep for it");
      return vtkSmartPointer<vtkDataObject>();
    }
    return rggAuxiliaryGeometryExtension::generateRGGDuctRepresentation(rggEntity, genNormals);
  }
  else if (rggEntity.stringProperty("rggType")[0] == OperationHelper::protoTypeDescription)
  {
    return rggAuxiliaryGeometryExtension::generateRGGProtoTypeRepresentation(rggEntity, genNormals);
  }
  return vtkSmartPointer<vtkDataObject>();
}

vtkSmartPointer<vtkDataObject> rggAuxiliaryGeometryExtension::generateRGGPinRepresentation(
  const AuxiliaryGeometry& pinAux, bool genNormals)
{
  rggAuxiliaryGeometryExtension::ensureCache();
  if (!pinAux.hasStringProperty(Pin::propDescription))
  {
    smtkErrorMacro(smtk::io::Logger::instance(), "Cannot create a representation for pin "
          << pinAux.name() << "without json representation.");
    return vtkSmartPointer<vtkDataObject>();
  }
  Pin pin = json::parse(pinAux.stringProperty(Pin::propDescription)[0]);

  //Extract info from pin
  int materialIndex = pin.cellMaterialIndex();
  bool isMaterialSet = materialIndex > 0 ? true : false;

  double zOrigin = pin.zOrigin();

  bool isHex = pinAux.owningModel().integerProperty(Core::geomDescription)[0]
                                      == static_cast<int>(Core::GeomType::Hex);

  int isCutAway = pin.isCutAway();

  const std::vector<Pin::Piece>& pieces = pin.pieces();

  const std::vector<Pin::LayerMaterial>& layerMs = pin.layerMaterials();

  // Follow logic in cmbNucRender::createGeo function. L249
  // Create a name-auxgeom map so that we can assign the right rep
  AuxiliaryGeometries childrenAux = pinAux.auxiliaryGeometries();
  std::map<std::string, AuxiliaryGeometry*> nameToChildAux;
  for (auto aux : childrenAux)
  {
    nameToChildAux[aux.name()] = &aux;
  }

  size_t numParts = pieces.size();
  size_t numLayers = layerMs.size();
  double baseCenter(zOrigin);

  constexpr int PinCellResolution = 20;

  for (size_t j = 0; j < numParts; ++j)
  {
    /// Create a vtkCmbLayeredConeSource for current part
    vtkSmartPointer<vtkCmbLayeredConeSource> coneSource =
      vtkSmartPointer<vtkCmbLayeredConeSource>::New();
    //  Add a material layer if needed
    coneSource->SetNumberOfLayers(static_cast<int>(numLayers + isMaterialSet));
    double height(pieces[j].length);
    double baseR(pieces[j].baseRadius), topR(pieces[j].topRadius);
    //    double layer42 = numLayers + isMaterialSet;
    //    std::cout << "Processing part " << j << " with base center as " <<
    //                 baseCenter << " ,height as " << height <<
    //                 " and number of layers to be " << layer42<<
    //                 " baseR=" << baseR << " topR=" << topR<<std::endl;
    // baseCenter would be updated at the end of the loop
    coneSource->SetBaseCenter(0, 0, baseCenter);
    coneSource->SetHeight(height);
    double largestRadius = 0;

    for (size_t k = 0; k < numLayers; k++)
    {
      // Calculate the baseR and topR at current layer
      double baseRL = baseR * layerMs[k].normRadius;
      double topRL = topR * layerMs[k].normRadius;
      coneSource->SetBaseRadius(static_cast<int>(k), baseRL);
      coneSource->SetTopRadius(static_cast<int>(k), topRL);
      coneSource->SetResolution(static_cast<int>(k), PinCellResolution);
      // Update largest radius for cell material visualization
      if (largestRadius < baseRL)
      {
        largestRadius = baseRL;
      }
      if (largestRadius < topRL)
      {
        largestRadius = topRL;
      }
    }
    if (isMaterialSet) // We have a valid material assigned( 0 means no material)
    {
      largestRadius *= 2.50;
      double r[] = { largestRadius * 0.5, largestRadius * 0.5 };
      int res = 4;
      if (isHex)
      {
        res = 6;
        r[0] = r[1] = r[0] / cos30;
      }
      coneSource->SetBaseRadius(static_cast<int>(numLayers), r[0], r[1]);
      coneSource->SetTopRadius(static_cast<int>(numLayers), r[0], r[1]);
      coneSource->SetResolution(static_cast<int>(numLayers), res);
    }
    double direction[] = { 0, 0, 1 };
    coneSource->SetDirection(direction);

    coneSource->SetGenerateNormals(genNormals);

    if (coneSource == nullptr)
    {
      continue;
    }
    // Cache child auxgeom(layer and part) with their polydata
    for (size_t k = 0; k < numLayers + isMaterialSet; k++)
    {
      std::string subName = pin.name() + smtk::session::rgg::Pin::subpartDescription + std::to_string(j) +
        rggPin::layerDescription + std::to_string(k);
      if (isMaterialSet && k == numLayers)
      {
        subName = pin.name() + smtk::session::rgg::Pin::subpartDescription + std::to_string(j) +
          rggPin::materialDescription;
      }
      // Follow logic in L263 cmbNucRender
      vtkSmartPointer<vtkPolyData> dataset = coneSource->CreateUnitLayer(static_cast<int>(k));
      // Since it's a unit layer, proper trasformation should be applied
      vtkSmartPointer<vtkTransform> transform = vtkSmartPointer<vtkTransform>::New();
      // Translate, rotate then scale
      double xyz[3];
      coneSource->GetBaseCenter(xyz);
      transform->Translate(xyz[0], xyz[1], xyz[2]);

      transform->RotateZ((isHex) ? 30 : 0);
      double angle = (isHex) ? 30 : 0;

      if (pieces[j].pieceType == rggPin::PieceType::CYLINDER && k == 0)
      { // Cylinder in the 0 layer should be handled differently(Following RGG's logic)
        transform->Scale(coneSource->GetTopRadius(static_cast<int>(k)),
          coneSource->GetBaseRadius(static_cast<int>(k)), height);
      }
      else
      {
        transform->Scale(1, 1, height);
      }
      vtkSmartPointer<vtkTransformPolyDataFilter> transformFilter =
        vtkSmartPointer<vtkTransformPolyDataFilter>::New();
      transformFilter->SetInputData(dataset);
      transformFilter->SetTransform(transform);
      transformFilter->Update();
      vtkSmartPointer<vtkPolyData> transformed = transformFilter->GetOutput();

      // Check cut away flag
      if (isCutAway)
      {
        double normal[] = { 0, 1, 0 };
        double xform[] = { 0, 0, -angle };

        transformNormal(normal, xform);
        clip(transformed, transformed, normal);
      }
      // Find the right auxgeom
      // FIXME: use the nameToChildAux map. For now if I deference the pointer
      // to a const ref, model builder would crash
      std::time_t mtime;
      std::time(&mtime);
      for (const auto& childAux : childrenAux)
      {
        if (childAux.name() == subName)
        {
          vtkAuxiliaryGeometryExtension::addCacheGeometry(transformed, childAux, mtime, false);
        }
      }
    }
    // Check if needed to create a boundary layer for the pin
    // Update current baseCenter
    baseCenter += height;
  }
  // Instead of return an aggregated multiblock, we generate a vtkObject for each sub auxgeom
  // due to a limitation of vtk that there is no easy way to specify colors for child blocks
  // in a nested vtk composite dataset. Anyway you have to flat it
  // out and specific color in a color array
  return vtkSmartPointer<vtkDataObject>();
}

vtkSmartPointer<vtkDataObject> rggAuxiliaryGeometryExtension::generateRGGDuctRepresentation(
  const AuxiliaryGeometry& ductAux, bool /*genNormals*/)
{
  rggAuxiliaryGeometryExtension::ensureCache();
  smtk::model::Model model = ductAux.owningModel();
  // Since duct properties are defined at core level, find the core from resource
  // then use its info to populate the panel.
  EntityRefArray coreArray =
      model.resource()->findEntitiesByProperty("rggType", Core::typeDescription);
  if (coreArray.size() != 1)
  {
    smtkErrorMacro(smtk::io::Logger::instance(), "Failed to find the core group");
    return vtkSmartPointer<vtkDataObject>();
  }
  EntityRef coreGroup = coreArray[0];
  Core core = json::parse(coreGroup.stringProperty(Core::propDescription)[0]);
  Duct duct = json::parse(ductAux.stringProperty(rggDuct::propDescription)[0]);

  bool isHex = core.geomType() == Core::GeomType::Hex;

  bool isCrossSection = duct.isCutAway();

  auto ductThickness = core.ductThickness();

  const auto& segments = duct.segments();

  // Follow logic in cmbNucRender::createGeo function. L168
  // Create a name-auxgeom map so that we can assign the right rep
  AuxiliaryGeometries childrenAux = ductAux.auxiliaryGeometries();
  std::map<std::string, AuxiliaryGeometry*> nameToChildAux;
  for (auto aux : childrenAux)
  {
    nameToChildAux[aux.name()] = &aux;
  }

  size_t numSegs = segments.size();
  for (size_t i = 0; i < numSegs; i++)
  {
    // Create layer manager
    // TODO: For now I just blindly follow the generation logic in RGG. It's not
    // straightforward and if we have time, we should simplify it.
    const auto& currentSeg = segments[i];
    double z1 =  currentSeg.baseZ;
    double height = currentSeg.height;
    double magicRatio = 0.0005; // Magic number used in rgg appliation code.
    double deltaZ = height * magicRatio;
    if (i == 0)
    {
      z1 = z1 + deltaZ;
      // if more than one duct, first duct height need to be reduced by deltaZ
      height = (numSegs > 1) ? height - deltaZ : height - 2 * deltaZ;
    }
    else if (i == (numSegs - 1)) //last duct
    {
      height -= 2 * deltaZ;
    }
    else
    {
      z1 += deltaZ;
    }
    size_t numLayers = currentSeg.layers.size();
    vtkSmartPointer<vtkCmbLayeredConeSource> coneSource =
      vtkSmartPointer<vtkCmbLayeredConeSource>::New();
    coneSource->SetNumberOfLayers(static_cast<int>(numLayers));
    coneSource->SetBaseCenter(0, 0, z1);
    double direction[] = { 0, 0, 1 };
    coneSource->SetDirection(direction);
    coneSource->SetHeight(height);

    int res = 4;
    double mult = 0.5;

    if (isHex)
    {
      res = 6;
      mult = 0.5 / cos30;
    }

    for (size_t k = 0; k < numLayers; k++)
    { // For each layer based on is hex or not,
      // it might have two different thicknesses
      const auto& layer = currentSeg.layers[k];
      double tx = std::get<1>(layer) * ductThickness.first *(1- magicRatio);
      double ty = std::get<2>(layer) * ductThickness.second  * (1 - magicRatio);
      coneSource->SetBaseRadius(static_cast<int>(k), tx * mult, ty * mult);
      coneSource->SetTopRadius(static_cast<int>(k), tx * mult, ty * mult);
      coneSource->SetResolution(static_cast<int>(k), res);
    }

    // Cache child auxgeom(layer and part) with their polydata
    for (size_t k = 0; k < numLayers; k++)
    {
      std::string subName = duct.name() + smtk::session::rgg::Duct::segmentDescription + std::to_string(i) +
        smtk::session::rgg::Duct::layerDescription + std::to_string(k);

      // Follow logic in L168 cmbNucRender
      vtkSmartPointer<vtkPolyData> dataset = coneSource->CreateUnitLayer(static_cast<int>(k));
      // Since it's a unit layer, proper trasformation should be applied
      vtkSmartPointer<vtkTransform> transform = vtkSmartPointer<vtkTransform>::New();
      // Translate, rotate then scale
      double xyz[3];
      coneSource->GetBaseCenter(xyz);
      transform->Translate(xyz[0], xyz[1], xyz[2]);

      if (k == 0)
      { // Cylinder in the 0 layer should be handled differently(Following RGG's logic)
        transform->Scale(coneSource->GetTopRadius(static_cast<int>(k), 0),
          coneSource->GetBaseRadius(static_cast<int>(k), 1), height);
      }
      else
      {
        transform->Scale(1, 1, height);
      }

      vtkSmartPointer<vtkTransformPolyDataFilter> transformFilter =
        vtkSmartPointer<vtkTransformPolyDataFilter>::New();
      transformFilter->SetInputData(dataset);
      transformFilter->SetTransform(transform);
      transformFilter->Update();
      vtkSmartPointer<vtkPolyData> transformed = transformFilter->GetOutput();

      // Check cut away flag
      if (isCrossSection)
      {
        double normal[] = { 0, 1, 0 };
        clip(transformed, transformed, normal);
      }
      // Find the right auxgeom
      // FIXME: use the nameToChildAux map. For now if I deference the pointer
      // to a const ref, model builder would crash
      std::time_t mtime;
      std::time(&mtime);
      for (const auto& childAux : childrenAux)
      {
        if (childAux.name() == subName)
        {
          vtkAuxiliaryGeometryExtension::addCacheGeometry(transformed, childAux, mtime, false);
        }
      }
    }
  }
  // Instead of return an aggregated multiblock, we generate a vtkObject for each sub auxgeom
  // due to a limitation of vtk that there is no easy way to specify colors for child blocks
  // in a nested vtk composite dataset. Anyway you have to flat it
  // out and specific color in a color array
  return vtkSmartPointer<vtkDataObject>();
}

vtkSmartPointer<vtkDataObject> rggAuxiliaryGeometryExtension::generateRGGProtoTypeRepresentation(
  const AuxiliaryGeometry& protoAux, bool /*genNormals*/)
{
  rggAuxiliaryGeometryExtension::ensureCache();
  if (!protoAux.hasStringProperty("representation_type"))
  {
    smtkErrorMacro(smtk::io::Logger::instance(), "Failed to find the representation_type for prototype "
                   << protoAux.name());
    return vtkSmartPointer<vtkDataObject>();
  }
  if (protoAux.stringProperty("representation_type")[0] == rggDuct::typeDescription)
  {
    smtk::model::Model model = protoAux.owningModel();
    // Since duct properties are defined at core level, find the core from resource
    // then use its info to populate the panel.
    EntityRefArray coreArray =
        model.resource()->findEntitiesByProperty("rggType", Core::typeDescription);
    if (coreArray.size() != 1)
    {
      smtkErrorMacro(smtk::io::Logger::instance(), "Failed to find the core group for prototype " << protoAux.name());
      return vtkSmartPointer<vtkDataObject>();
    }
    EntityRef coreGroup = coreArray[0];
    Core core = json::parse(coreGroup.stringProperty(Core::propDescription)[0]);
    Duct duct = json::parse(protoAux.stringProperty("representation")[0]);
    smtk::model::IntegerList indexes = protoAux.integerProperty("index");
    size_t segIndex = static_cast<size_t>(indexes[0]);
    size_t layerIndex = static_cast<size_t>(indexes[1]);

    bool isHex = core.geomType() == Core::GeomType::Hex;

    auto ductThickness = core.ductThickness();

    const auto& segments = duct.segments();

    size_t numSegs = segments.size();
    // Create layer manager
    // TODO: For now I just blindly follow the generation logic in RGG. It's not
    // straightforward and if we have time, we should simplify it.
    const auto& currentSeg = segments[segIndex];
    double z1 =  currentSeg.baseZ;
    double height = currentSeg.height;
    double magicRatio = 0.0005; // Magic number used in rgg appliation code.
    double deltaZ = height * magicRatio;
    if (segIndex == 0)
    {
      z1 = z1 + deltaZ;
      // if more than one duct, first duct height need to be reduced by deltaZ
      height = (numSegs > 1) ? height - deltaZ : height - 2 * deltaZ;
    }
    else if (segIndex == (numSegs - 1)) //last duct
    {
      height -= 2 * deltaZ;
    }
    else
    {
      z1 += deltaZ;
    }
    size_t numLayers = currentSeg.layers.size();
    vtkSmartPointer<vtkCmbLayeredConeSource> coneSource =
      vtkSmartPointer<vtkCmbLayeredConeSource>::New();
    coneSource->SetNumberOfLayers(static_cast<int>(numLayers));
    coneSource->SetBaseCenter(0, 0, z1);
    double direction[] = { 0, 0, 1 };
    coneSource->SetDirection(direction);
    coneSource->SetHeight(height);

    int res = 4;
    double mult = 0.5;

    if (isHex)
    {
      res = 6;
      mult = 0.5 / cos30;
    }

    for (size_t k = 0; k < numLayers; k++)
    { // For each layer based on is hex or not,
      // it might have two different thicknesses
      const auto& layer = currentSeg.layers[k];
      double tx = std::get<1>(layer) * ductThickness.first *(1- magicRatio);
      double ty = std::get<2>(layer) * ductThickness.second  * (1 - magicRatio);
      coneSource->SetBaseRadius(static_cast<int>(k), tx * mult, ty * mult);
      coneSource->SetTopRadius(static_cast<int>(k), tx * mult, ty * mult);
      coneSource->SetResolution(static_cast<int>(k), res);
    }

    // Follow logic in L168 cmbNucRender
    vtkSmartPointer<vtkPolyData> dataset = coneSource->CreateUnitLayer(static_cast<int>(layerIndex));
    return dataset;
  }
  else if (protoAux.stringProperty("representation_type")[0] == rggPin::typeDescription)
  {
    Pin pin = json::parse(protoAux.stringProperty("representation")[0]);
    smtk::model::IntegerList indexes = protoAux.integerProperty("index");
    size_t pieceIndex = static_cast<size_t>(indexes[0]);
    size_t layerIndex = static_cast<size_t>(indexes[1]);

    //Extract info from pin
    int materialIndex = pin.cellMaterialIndex();
    bool isMaterialSet = materialIndex > 0 ? true : false;

    double zOrigin = pin.zOrigin();

    bool isHex = protoAux.owningModel().integerProperty(Core::geomDescription)[0]
                                        == static_cast<int>(Core::GeomType::Hex);

    const std::vector<Pin::Piece>& pieces = pin.pieces();

    const std::vector<Pin::LayerMaterial>& layerMs = pin.layerMaterials();

    size_t numLayers = layerMs.size();
    double baseCenter(zOrigin);

    constexpr int PinCellResolution = 20;

    // Update current baseCenter
    for (size_t partI = 0; partI < pieceIndex; ++partI)
    {
      double height(pieces[pieceIndex].length);
      baseCenter += height;
    }
      /// Create a vtkCmbLayeredConeSource for current part
      vtkSmartPointer<vtkCmbLayeredConeSource> coneSource =
        vtkSmartPointer<vtkCmbLayeredConeSource>::New();
      //  Add a material layer if needed
      coneSource->SetNumberOfLayers(static_cast<int>(numLayers + isMaterialSet));
      double height(pieces[pieceIndex].length);
      double baseR(pieces[pieceIndex].baseRadius), topR(pieces[pieceIndex].topRadius);
      // baseCenter would be updated at the end of the loop
      coneSource->SetBaseCenter(0, 0, baseCenter);
      coneSource->SetHeight(height);
      double largestRadius = 0;

      for (size_t k = 0; k < numLayers; k++)
      {
        // Calculate the baseR and topR at current layer
        double baseRL = baseR * layerMs[k].normRadius;
        double topRL = topR * layerMs[k].normRadius;
        coneSource->SetBaseRadius(static_cast<int>(k), baseRL);
        coneSource->SetTopRadius(static_cast<int>(k), topRL);
        coneSource->SetResolution(static_cast<int>(k), PinCellResolution);
        // Update largest radius for cell material visualization
        if (largestRadius < baseRL)
        {
          largestRadius = baseRL;
        }
        if (largestRadius < topRL)
        {
          largestRadius = topRL;
        }
      }
      if (isMaterialSet) // We have a valid material assigned( 0 means no material)
      {
        largestRadius *= 2.50;
        double r[] = { largestRadius * 0.5, largestRadius * 0.5 };
        int res = 4;
        if (isHex)
        {
          res = 6;
          r[0] = r[1] = r[0] / cos30;
        }
        coneSource->SetBaseRadius(static_cast<int>(numLayers), r[0], r[1]);
        coneSource->SetTopRadius(static_cast<int>(numLayers), r[0], r[1]);
        coneSource->SetResolution(static_cast<int>(numLayers), res);
      }
      double direction[] = { 0, 0, 1 };
      coneSource->SetDirection(direction);

      //coneSource->SetGenerateNormals(genNormals);

      // Follow logic in L263 cmbNucRender
      vtkSmartPointer<vtkPolyData> dataset = coneSource->CreateUnitLayer(static_cast<int>(layerIndex));
      return dataset;
  }
  else
  {
    smtkErrorMacro(smtk::io::Logger::instance(), "Invalid representation_type for prototype "
                   << protoAux.name());
    return vtkSmartPointer<vtkDataObject>();
  }
}

smtkDeclareExtension(
  /*SMTKRGGSESSION_EXPORT*/, rgg_auxiliary_geometry, rggAuxiliaryGeometryExtension);
smtkComponentInitMacro(smtk_rgg_auxiliary_geometry_extension);
