//============================================================================
//  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.
//
//  Copyright 2015 Sandia Corporation.
//  Copyright 2015 UT-Battelle, LLC.
//  Copyright 2015 Los Alamos National Security.
//
//  Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
//  the U.S. Government retains certain rights in this software.
//
//  Under the terms of Contract DE-AC52-06NA25396 with Los Alamos National
//  Laboratory (LANL), the U.S. Government retains certain rights in
//  this software.
//============================================================================
#ifndef vtk_m_rendering_Triangulator_h
#define vtk_m_rendering_Triangulator_h

#include <vtkm/cont/ArrayHandleCounting.h>
#include <vtkm/cont/DataSet.h>
#include <vtkm/exec/ExecutionWholeArray.h>
#include <vtkm/worklet/DispatcherMapField.h>
#include <vtkm/worklet/DispatcherMapTopology.h>
#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletMapTopology.h>

namespace vtkm
{
namespace rendering
{
/// \brief Triangulator creates a minimal set of triangles from a cell set.
///
///  This class creates a array of triangle indices from both 3D and 2D
///  explicit cell sets. This list can serve as input to opengl and the
///  ray tracer scene renderers. TODO: Add regular grid support
///
template <typename DeviceAdapter>
class Triangulator
{
private:
  typedef typename vtkm::cont::ArrayHandle<vtkm::Id> IdArrayHandle;
  typedef typename vtkm::cont::ArrayHandle<vtkm::Vec<vtkm::Id, 4>> Vec4ArrayHandle;
  typedef typename Vec4ArrayHandle::ExecutionTypes<DeviceAdapter>::Portal Vec4ArrayPortalType;
  typedef typename IdArrayHandle::ExecutionTypes<DeviceAdapter>::PortalConst IdPortalConstType;

public:
  template <class T>
  class MemSet : public vtkm::worklet::WorkletMapField
  {
    T Value;

  public:
    VTKM_CONT
    MemSet(T value)
      : Value(value)
    {
    }
    typedef void ControlSignature(FieldOut<>);
    typedef void ExecutionSignature(_1);
    VTKM_EXEC
    void operator()(T& outValue) const { outValue = Value; }
  }; //class MemSet
  class CountTriangles : public vtkm::worklet::WorkletMapField
  {
  public:
    VTKM_CONT
    CountTriangles() {}
    typedef void ControlSignature(FieldIn<>, FieldOut<>);
    typedef void ExecutionSignature(_1, _2);
    VTKM_EXEC
    void operator()(const vtkm::Id& shapeType, vtkm::Id& triangles) const
    {
      if (shapeType == vtkm::CELL_SHAPE_TRIANGLE)
        triangles = 1;
      else if (shapeType == vtkm::CELL_SHAPE_QUAD)
        triangles = 2;
      else if (shapeType == vtkm::CELL_SHAPE_TETRA)
        triangles = 4;
      else if (shapeType == vtkm::CELL_SHAPE_HEXAHEDRON)
        triangles = 12;
      else if (shapeType == vtkm::CELL_SHAPE_WEDGE)
        triangles = 8;
      else if (shapeType == vtkm::CELL_SHAPE_PYRAMID)
        triangles = 6;
      else
        triangles = 0;
    }
  }; //class CountTriangles

  template <int DIM>
  class TrianglulateStructured : public vtkm::worklet::WorkletMapPointToCell
  {
  private:
    Vec4ArrayPortalType OutputIndices;

  public:
    typedef void ControlSignature(CellSetIn cellset, FieldInTo<>);
    typedef void ExecutionSignature(FromIndices, _2);
    //typedef _1 InputDomain;
    VTKM_CONT
    TrianglulateStructured(vtkm::cont::ArrayHandle<vtkm::Vec<vtkm::Id, 4>>& outputIndices)
    {
      this->OutputIndices =
        outputIndices.PrepareForOutput(outputIndices.GetNumberOfValues(), DeviceAdapter());
    }

    //TODO: Remove the if/then with templates.
    template <typename CellNodeVecType>
    VTKM_EXEC void operator()(const CellNodeVecType& cellIndices, const vtkm::Id& cellIndex) const
    {
      vtkm::Vec<vtkm::Id, 4> triangle;
      if (DIM == 2)
      {
        const vtkm::Id triangleOffset = cellIndex * 2;
        // 0-1-2
        triangle[1] = cellIndices[0];
        triangle[2] = cellIndices[1];
        triangle[3] = cellIndices[2];
        triangle[0] = cellIndex;
        OutputIndices.Set(triangleOffset, triangle);
        // 0-3-2
        triangle[2] = cellIndices[3];
        OutputIndices.Set(triangleOffset + 1, triangle);
      }
      else if (DIM == 3)
      {
        const vtkm::Id triangleOffset = cellIndex * 12;
        // 0-1-2
        triangle[1] = cellIndices[0];
        triangle[2] = cellIndices[1];
        triangle[3] = cellIndices[2];
        triangle[0] = cellIndex;
        OutputIndices.Set(triangleOffset, triangle);
        // 0-3-2
        triangle[2] = cellIndices[3];
        OutputIndices.Set(triangleOffset + 1, triangle);
        // 0-3-7
        triangle[3] = cellIndices[7];
        OutputIndices.Set(triangleOffset + 2, triangle);
        // 0-4-7
        triangle[2] = cellIndices[4];
        OutputIndices.Set(triangleOffset + 3, triangle);
        // 5-4-7
        triangle[1] = cellIndices[5];
        OutputIndices.Set(triangleOffset + 4, triangle);
        // 5-6-7
        triangle[2] = cellIndices[6];
        OutputIndices.Set(triangleOffset + 5, triangle);
        // 3-6-7
        triangle[1] = cellIndices[3];
        OutputIndices.Set(triangleOffset + 6, triangle);
        // 3-6-2
        triangle[3] = cellIndices[2];
        OutputIndices.Set(triangleOffset + 7, triangle);
        // 1-6-2
        triangle[1] = cellIndices[1];
        OutputIndices.Set(triangleOffset + 8, triangle);
        // 1-6-5
        triangle[3] = cellIndices[5];
        OutputIndices.Set(triangleOffset + 9, triangle);
        // 1-4-5
        triangle[2] = cellIndices[4];
        OutputIndices.Set(triangleOffset + 10, triangle);
        // 1-4-0
        triangle[3] = cellIndices[0];
        OutputIndices.Set(triangleOffset + 11, triangle);
      }
    }
  };

  class IndicesSort : public vtkm::worklet::WorkletMapField
  {
  public:
    VTKM_CONT
    IndicesSort() {}
    typedef void ControlSignature(FieldInOut<>);
    typedef void ExecutionSignature(_1);
    VTKM_EXEC
    void operator()(vtkm::Vec<vtkm::Id, 4>& triangleIndices) const
    {
      // first field contains the id of the cell the
      // trianlge belongs to
      vtkm::Id temp;
      if (triangleIndices[1] > triangleIndices[3])
      {
        temp = triangleIndices[1];
        triangleIndices[1] = triangleIndices[3];
        triangleIndices[3] = temp;
      }
      if (triangleIndices[1] > triangleIndices[2])
      {
        temp = triangleIndices[1];
        triangleIndices[1] = triangleIndices[2];
        triangleIndices[2] = temp;
      }
      if (triangleIndices[2] > triangleIndices[3])
      {
        temp = triangleIndices[2];
        triangleIndices[2] = triangleIndices[3];
        triangleIndices[3] = temp;
      }
    }
  }; //class IndicesSort

  struct IndicesLessThan
  {
    VTKM_EXEC_CONT
    bool operator()(const vtkm::Vec<vtkm::Id, 4>& a, const vtkm::Vec<vtkm::Id, 4>& b) const
    {
      if (a[1] < b[1])
        return true;
      if (a[1] > b[1])
        return false;
      if (a[2] < b[2])
        return true;
      if (a[2] > b[2])
        return false;
      if (a[3] < b[3])
        return true;
      return false;
    }
  };

  class UniqueTriangles : public vtkm::worklet::WorkletMapField
  {
  public:
    VTKM_CONT
    UniqueTriangles() {}
    typedef void ControlSignature(ExecObject, ExecObject);
    typedef void ExecutionSignature(_1, _2, WorkIndex);
    VTKM_EXEC
    bool IsTwin(const vtkm::Vec<vtkm::Id, 4>& a, const vtkm::Vec<vtkm::Id, 4>& b) const
    {
      return (a[1] == b[1] && a[2] == b[2] && a[3] == b[3]);
    }
    VTKM_EXEC
    void operator()(vtkm::exec::ExecutionWholeArrayConst<vtkm::Vec<vtkm::Id, 4>>& indices,
                    vtkm::exec::ExecutionWholeArray<vtkm::UInt8>& outputFlags,
                    const vtkm::Id& index) const
    {
      if (index == 0)
        return;
      //if we are a shared face, mark ourself and neighbor for desctruction
      if (IsTwin(indices.Get(index), indices.Get(index - 1)))
      {
        outputFlags.Set(index, 0);
        outputFlags.Set(index - 1, 0);
      }
    }
  }; //class UniqueTriangles

  class Trianglulate : public vtkm::worklet::WorkletMapField
  {
  private:
    IdPortalConstType Indices;
    Vec4ArrayPortalType OutputIndices;

    vtkm::Int32 TRIANGLE_INDICES;
    vtkm::Int32 QUAD_INDICES;
    vtkm::Int32 TETRA_INDICES;
    vtkm::Int32 HEX_INDICES;
    vtkm::Int32 WEDGE_INDICES;
    vtkm::Int32 PYRAMID_INDICES;

  public:
    VTKM_CONT
    Trianglulate(vtkm::cont::ArrayHandle<vtkm::Vec<vtkm::Id, 4>>& outputIndices,
                 const vtkm::cont::ArrayHandle<vtkm::Id>& indices, const vtkm::Id& size)
      : Indices(indices.PrepareForInput(DeviceAdapter()))
    {
      this->OutputIndices = outputIndices.PrepareForOutput(size, DeviceAdapter());
      TRIANGLE_INDICES = 3;
      QUAD_INDICES = 4;
      TETRA_INDICES = 4;
      HEX_INDICES = 8;
      WEDGE_INDICES = 6;
      PYRAMID_INDICES = 5;
    }
    typedef void ControlSignature(FieldIn<>, FieldIn<>, FieldIn<>);
    typedef void ExecutionSignature(_1, _2, _3, WorkIndex);
    VTKM_EXEC
    void operator()(const vtkm::Id& shapeType, const vtkm::Id& indexOffset,
                    const vtkm::Id& triangleOffset, const vtkm::Id& cellId) const
    {
      vtkm::Vec<vtkm::Id, 4> triangle;
      if (shapeType == vtkm::CELL_SHAPE_TRIANGLE)
      {

        triangle[1] = Indices.Get(indexOffset + 0);
        triangle[2] = Indices.Get(indexOffset + 1);
        triangle[3] = Indices.Get(indexOffset + 2);
        triangle[0] = cellId;
        OutputIndices.Set(triangleOffset, triangle);
      }
      if (shapeType == vtkm::CELL_SHAPE_QUAD)
      {

        triangle[1] = Indices.Get(indexOffset + 0);
        triangle[2] = Indices.Get(indexOffset + 1);
        triangle[3] = Indices.Get(indexOffset + 2);
        triangle[0] = cellId;
        OutputIndices.Set(triangleOffset, triangle);

        triangle[2] = Indices.Get(indexOffset + 3);
        OutputIndices.Set(triangleOffset + 1, triangle);
      }
      if (shapeType == vtkm::CELL_SHAPE_TETRA)
      {
        // 0-1-2
        triangle[1] = Indices.Get(indexOffset + 1);
        triangle[2] = Indices.Get(indexOffset + 0);
        triangle[3] = Indices.Get(indexOffset + 2);
        triangle[0] = cellId;
        OutputIndices.Set(triangleOffset, triangle);
        // 0-1-3
        triangle[1] = Indices.Get(indexOffset + 1);
        triangle[2] = Indices.Get(indexOffset + 0);
        triangle[3] = Indices.Get(indexOffset + 3);

        OutputIndices.Set(triangleOffset + 1, triangle);
        // 2-1-3
        triangle[1] = Indices.Get(indexOffset + 1);
        triangle[2] = Indices.Get(indexOffset + 2);
        triangle[3] = Indices.Get(indexOffset + 3);
        OutputIndices.Set(triangleOffset + 2, triangle);
        // 2-0-3
        triangle[1] = Indices.Get(indexOffset + 2);
        triangle[2] = Indices.Get(indexOffset + 0);
        triangle[3] = Indices.Get(indexOffset + 3);
        OutputIndices.Set(triangleOffset + 3, triangle);
      }
      if (shapeType == vtkm::CELL_SHAPE_HEXAHEDRON)
      {
        // 0-1-2
        triangle[1] = Indices.Get(indexOffset + 0);
        triangle[2] = Indices.Get(indexOffset + 1);
        triangle[3] = Indices.Get(indexOffset + 2);
        triangle[0] = cellId;
        OutputIndices.Set(triangleOffset, triangle);
        // 0-3-2
        triangle[2] = Indices.Get(indexOffset + 3);
        OutputIndices.Set(triangleOffset + 1, triangle);
        // 0-3-7
        triangle[3] = Indices.Get(indexOffset + 7);
        OutputIndices.Set(triangleOffset + 2, triangle);
        // 0-4-7
        triangle[2] = Indices.Get(indexOffset + 4);
        OutputIndices.Set(triangleOffset + 3, triangle);
        // 5-4-7
        triangle[1] = Indices.Get(indexOffset + 5);
        OutputIndices.Set(triangleOffset + 4, triangle);
        // 5-6-7
        triangle[2] = Indices.Get(indexOffset + 6);
        OutputIndices.Set(triangleOffset + 5, triangle);
        // 3-6-7
        triangle[1] = Indices.Get(indexOffset + 3);
        OutputIndices.Set(triangleOffset + 6, triangle);
        // 3-6-2
        triangle[3] = Indices.Get(indexOffset + 2);
        OutputIndices.Set(triangleOffset + 7, triangle);
        // 1-6-2
        triangle[1] = Indices.Get(indexOffset + 1);
        OutputIndices.Set(triangleOffset + 8, triangle);
        // 1-6-5
        triangle[3] = Indices.Get(indexOffset + 5);
        OutputIndices.Set(triangleOffset + 9, triangle);
        // 1-4-5
        triangle[2] = Indices.Get(indexOffset + 4);
        OutputIndices.Set(triangleOffset + 10, triangle);
        // 1-4-0
        triangle[3] = Indices.Get(indexOffset + 0);
        OutputIndices.Set(triangleOffset + 11, triangle);
      }
      if (shapeType == vtkm::CELL_SHAPE_WEDGE)
      {
        // 0-1-2
        triangle[1] = Indices.Get(indexOffset + 0);
        triangle[2] = Indices.Get(indexOffset + 1);
        triangle[3] = Indices.Get(indexOffset + 2);
        triangle[0] = cellId;
        OutputIndices.Set(triangleOffset, triangle);
        // 0-3-2
        triangle[2] = Indices.Get(indexOffset + 3);
        OutputIndices.Set(triangleOffset + 1, triangle);
        // 5-3-2
        triangle[1] = Indices.Get(indexOffset + 5);
        OutputIndices.Set(triangleOffset + 2, triangle);
        // 5-3-4
        triangle[3] = Indices.Get(indexOffset + 4);
        OutputIndices.Set(triangleOffset + 3, triangle);
        // 5-2-4
        triangle[2] = Indices.Get(indexOffset + 2);
        OutputIndices.Set(triangleOffset + 4, triangle);
        // 5-1-4
        triangle[2] = Indices.Get(indexOffset + 1);
        OutputIndices.Set(triangleOffset + 5, triangle);
      }
      if (shapeType == vtkm::CELL_SHAPE_PYRAMID)
      {
        // 0-1-2
        triangle[1] = Indices.Get(indexOffset + 0);
        triangle[2] = Indices.Get(indexOffset + 1);
        triangle[3] = Indices.Get(indexOffset + 2);
        triangle[0] = cellId;
        OutputIndices.Set(triangleOffset, triangle);
        // 0-3-2
        triangle[2] = Indices.Get(indexOffset + 3);
        OutputIndices.Set(triangleOffset + 1, triangle);
        // 0-3-4
        triangle[3] = Indices.Get(indexOffset + 4);
        OutputIndices.Set(triangleOffset + 2, triangle);
        // 2-3-4
        triangle[1] = Indices.Get(indexOffset + 2);
        OutputIndices.Set(triangleOffset + 3, triangle);
        // 2-1-4
        triangle[2] = Indices.Get(indexOffset + 1);
        OutputIndices.Set(triangleOffset + 4, triangle);
        // 0-3-4
        triangle[1] = Indices.Get(indexOffset + 0);
        OutputIndices.Set(triangleOffset + 5, triangle);
      }
    }
  }; //class Trianglulate

public:
  VTKM_CONT
  Triangulator() {}

  VTKM_CONT
  void ExternalTrianlges(vtkm::cont::ArrayHandle<vtkm::Vec<vtkm::Id, 4>>& outputIndices,
                         vtkm::Id& outputTriangles)
  {
    //Eliminate unseen triangles
    vtkm::worklet::DispatcherMapField<IndicesSort>().Invoke(outputIndices);
    vtkm::cont::DeviceAdapterAlgorithm<DeviceAdapter>::Sort(outputIndices, IndicesLessThan());
    vtkm::cont::ArrayHandle<vtkm::UInt8> flags;
    flags.Allocate(outputTriangles);
    vtkm::worklet::DispatcherMapField<MemSet<vtkm::UInt8>>(MemSet<vtkm::UInt8>(1)).Invoke(flags);
    //Unique triangles will have a flag = 1
    vtkm::worklet::DispatcherMapField<UniqueTriangles>().Invoke(
      vtkm::exec::ExecutionWholeArrayConst<vtkm::Vec<vtkm::Id, 4>>(outputIndices),
      vtkm::exec::ExecutionWholeArray<vtkm::UInt8>(flags));

    vtkm::cont::ArrayHandle<vtkm::Vec<vtkm::Id, 4>> subset;
    vtkm::cont::DeviceAdapterAlgorithm<DeviceAdapter>::CopyIf(outputIndices, flags, subset);
    outputIndices = subset;
    outputTriangles = subset.GetNumberOfValues();
  }

  VTKM_CONT
  void Run(const vtkm::cont::DynamicCellSet& cellset,
           vtkm::cont::ArrayHandle<vtkm::Vec<vtkm::Id, 4>>& outputIndices,
           vtkm::Id& outputTriangles)
  {

    if (cellset.IsSameType(vtkm::cont::CellSetStructured<3>()))
    {
      vtkm::cont::CellSetStructured<3> cellSetStructured3D =
        cellset.Cast<vtkm::cont::CellSetStructured<3>>();
      const vtkm::Id numCells = cellSetStructured3D.GetNumberOfCells();

      vtkm::cont::ArrayHandleCounting<vtkm::Id> cellIdxs(0, 1, numCells);
      outputIndices.Allocate(numCells * 12);
      vtkm::worklet::DispatcherMapTopology<TrianglulateStructured<3>>(
        TrianglulateStructured<3>(outputIndices))
        .Invoke(cellSetStructured3D, cellIdxs);

      outputTriangles = numCells * 12;
    }
    else if (cellset.IsSameType(vtkm::cont::CellSetStructured<2>()))
    {
      vtkm::cont::CellSetStructured<2> cellSetStructured2D =
        cellset.Cast<vtkm::cont::CellSetStructured<2>>();
      const vtkm::Id numCells = cellSetStructured2D.GetNumberOfCells();

      vtkm::cont::ArrayHandleCounting<vtkm::Id> cellIdxs(0, 1, numCells);
      outputIndices.Allocate(numCells * 2);
      vtkm::worklet::DispatcherMapTopology<TrianglulateStructured<2>>(
        TrianglulateStructured<2>(outputIndices))
        .Invoke(cellSetStructured2D, cellIdxs);

      outputTriangles = numCells * 2;
    }
    else if (cellset.IsSameType(vtkm::cont::CellSetExplicit<>()))
    {
      vtkm::cont::CellSetExplicit<> cellSetExplicit = cellset.Cast<vtkm::cont::CellSetExplicit<>>();
      const vtkm::cont::ArrayHandle<vtkm::UInt8> shapes = cellSetExplicit.GetShapesArray(
        vtkm::TopologyElementTagPoint(), vtkm::TopologyElementTagCell());
      const vtkm::cont::ArrayHandle<vtkm::Int32> indices = cellSetExplicit.GetNumIndicesArray(
        vtkm::TopologyElementTagPoint(), vtkm::TopologyElementTagCell());
      vtkm::cont::ArrayHandle<vtkm::Id> conn = cellSetExplicit.GetConnectivityArray(
        vtkm::TopologyElementTagPoint(), vtkm::TopologyElementTagCell());
      vtkm::cont::ArrayHandle<vtkm::Id> offsets = cellSetExplicit.GetIndexOffsetArray(
        vtkm::TopologyElementTagPoint(), vtkm::TopologyElementTagCell());

      // We need to somehow force the data set to build the index offsets
      //vtkm::IdComponent c = indices.GetPortalControl().Get(0);
      vtkm::Vec<vtkm::Id, 3> forceBuildIndices;
      cellSetExplicit.GetIndices(0, forceBuildIndices);

      vtkm::cont::ArrayHandle<vtkm::Id> trianglesPerCell;
      vtkm::worklet::DispatcherMapField<CountTriangles>(CountTriangles())
        .Invoke(shapes, trianglesPerCell);

      vtkm::Id totalTriangles = 0;
      totalTriangles =
        vtkm::cont::DeviceAdapterAlgorithm<DeviceAdapter>::Reduce(trianglesPerCell, vtkm::Id(0));

      vtkm::cont::ArrayHandle<vtkm::Id> cellOffsets;
      vtkm::cont::DeviceAdapterAlgorithm<DeviceAdapter>::ScanExclusive(trianglesPerCell,
                                                                       cellOffsets);
      outputIndices.Allocate(totalTriangles);
      vtkm::worklet::DispatcherMapField<Trianglulate>(
        Trianglulate(outputIndices, conn, totalTriangles))
        .Invoke(shapes, offsets, cellOffsets);

      outputTriangles = totalTriangles;
    }
    else if (cellset.IsSameType(vtkm::cont::CellSetSingleType<>()))
    {
      typedef vtkm::TopologyElementTagPoint PointTag;
      typedef vtkm::TopologyElementTagCell CellTag;

      vtkm::cont::CellSetSingleType<> cellSetSingleType =
        cellset.Cast<vtkm::cont::CellSetSingleType<>>();

      //fetch and see if we are all triangles, that currently is the only
      //cell set single type we support.
      vtkm::Id shapeTypeAsId = cellSetSingleType.GetCellShape(0);

      if (shapeTypeAsId == vtkm::CellShapeTagTriangle::Id)
      {
        //generate the outputIndices
        vtkm::Id totalTriangles = cellSetSingleType.GetNumberOfCells();
        vtkm::cont::ArrayHandleCounting<vtkm::Id> cellIdxs(0, 1, totalTriangles);

        outputIndices.Allocate(totalTriangles);
        vtkm::worklet::DispatcherMapField<Trianglulate>(
          Trianglulate(outputIndices, cellSetSingleType.GetConnectivityArray(PointTag(), CellTag()),
                       totalTriangles))
          .Invoke(cellSetSingleType.GetShapesArray(PointTag(), CellTag()),
                  cellSetSingleType.GetIndexOffsetArray(PointTag(), CellTag()), cellIdxs);

        outputTriangles = totalTriangles;
      }
      else
      {
        throw vtkm::cont::ErrorBadType(
          "Unsupported cell type for trianglulation with CellSetSingleType");
      }
    }
    else
    {
      throw vtkm::cont::ErrorBadType("Unsupported cell set type for trianglulation");
    }

    //get rid of any triagles we cannot see
    ExternalTrianlges(outputIndices, outputTriangles);
  }
}; // class Triangulator
}
} //namespace vtkm::rendering
#endif //vtk_m_rendering_Triangulator_h