// Created on: 2014-09-11
// Created by: Danila ULYANOV
// Copyright (c) 2013-2014 OPEN CASCADE SAS
//
// This file is part of Open CASCADE Technology software library.
//
// This library is free software; you can redistribute it and/or modify it under
// the terms of the GNU Lesser General Public License version 2.1 as published
// by the Free Software Foundation, with special exception defined in the file
// OCCT_LGPL_EXCEPTION.txt. Consult the file LICENSE_LGPL_21.txt included in OCCT
// distribution for complete text of the license and disclaimer of any warranty.
//
// Alternatively, this file may be used under the terms of Open CASCADE
// commercial license or contractual agreement.

#include <algorithm>

#include <Standard_Assert.hxx>

#include <NCollection_Array1.hxx>

#ifdef HAVE_TBB
  // On Windows, function TryEnterCriticalSection has appeared in Windows NT
  // and is surrounded by #ifdef in MS VC++ 7.1 headers.
  // Thus to use it we need to define appropriate macro saying that we will
  // run on Windows NT 4.0 at least
  #if defined(_WIN32) && !defined(_WIN32_WINNT)
    #define _WIN32_WINNT 0x0501
  #endif

  #include <tbb/task.h>
#endif

// =======================================================================
// function : BVH_LinearBuilder
// purpose  :
// =======================================================================
template<class T, int N>
BVH_LinearBuilder<T, N>::BVH_LinearBuilder (const Standard_Integer theLeafNodeSize,
                                            const Standard_Integer theMaxTreeDepth)
: BVH_Builder<T, N> (theLeafNodeSize,
                     theMaxTreeDepth)
{
  //
}

// =======================================================================
// function : ~BVH_LinearBuilder
// purpose  :
// =======================================================================
template<class T, int N>
BVH_LinearBuilder<T, N>::~BVH_LinearBuilder()
{
  //
}

namespace BVH
{
  // Radix sort STL predicate for 32-bit integer.
  class BitPredicate
  {
    Standard_Integer myBit;

  public:
    
    //! Creates new radix sort predicate.
    BitPredicate (const Standard_Integer theBit) : myBit (theBit) 
    {
      //
    }
    
    //! Returns predicate value.
    bool operator() (const BVH_EncodedLink theLink) const
    {
      const Standard_Integer aMask = 1 << myBit;

      return !(theLink.first & aMask); // 0-bit to the left side
    }
  };

  //! STL compare tool used in binary search algorithm.
  class BitComparator
  {
    Standard_Integer myBit;

  public:

    //! Creates new STL comparator.
    BitComparator (const Standard_Integer theBit) : myBit (theBit) 
    {
      //
    }

    //! Checks left value for the given bit.
    bool operator() (BVH_EncodedLink theLink1, BVH_EncodedLink /*theLink2*/)
    {
      return !(theLink1.first & (1 << myBit));
    }
  };

  //! Tool object for sorting link array using radix sort algorithm.
  struct RadixSorter
  {
    typedef std::vector<BVH_EncodedLink>::iterator LinkIterator;

    // Performs MSD (most significant digit) radix sort.
    static void Perform (LinkIterator theStart, LinkIterator theFinal, Standard_Integer theBit = 29)
    {
      while (theStart != theFinal && theBit >= 0)
      {
        LinkIterator anOffset = std::partition (theStart, theFinal, BitPredicate (theBit--));
        
        Perform (theStart, anOffset, theBit);
      
        theStart = anOffset;
      }
    }
  };

  //! Calculates bounding boxes (AABBs) for the given BVH tree. 
  template<class T, int N>
  Standard_Integer UpdateBounds (BVH_Set<T, N>* theSet, BVH_Tree<T, N>* theTree, const Standard_Integer theNode = 0)
  {
    const BVH_Vec4i aData = theTree->NodeInfoBuffer()[theNode];

    if (aData.x() == 0)
    {
      const Standard_Integer aLftChild = theTree->NodeInfoBuffer()[theNode].y();
      const Standard_Integer aRghChild = theTree->NodeInfoBuffer()[theNode].z();

      const Standard_Integer aLftDepth = UpdateBounds (theSet, theTree, aLftChild);
      const Standard_Integer aRghDepth = UpdateBounds (theSet, theTree, aRghChild);

      typename BVH_Box<T, N>::BVH_VecNt aLftMinPoint = theTree->MinPointBuffer()[aLftChild];
      typename BVH_Box<T, N>::BVH_VecNt aLftMaxPoint = theTree->MaxPointBuffer()[aLftChild];
      typename BVH_Box<T, N>::BVH_VecNt aRghMinPoint = theTree->MinPointBuffer()[aRghChild];
      typename BVH_Box<T, N>::BVH_VecNt aRghMaxPoint = theTree->MaxPointBuffer()[aRghChild];

      BVH::BoxMinMax<T, N>::CwiseMin (aLftMinPoint, aRghMinPoint);
      BVH::BoxMinMax<T, N>::CwiseMax (aLftMaxPoint, aRghMaxPoint);
    
      theTree->MinPointBuffer()[theNode] = aLftMinPoint;
      theTree->MaxPointBuffer()[theNode] = aLftMaxPoint;

      return Max (aLftDepth, aRghDepth) + 1;
    }
    else
    {
      typename BVH_Box<T, N>::BVH_VecNt& aMinPoint = theTree->MinPointBuffer()[theNode];
      typename BVH_Box<T, N>::BVH_VecNt& aMaxPoint = theTree->MaxPointBuffer()[theNode];

      for (Standard_Integer aPrimIdx = aData.y(); aPrimIdx <= aData.z(); ++aPrimIdx)
      {
        const BVH_Box<T, N> aBox = theSet->Box (aPrimIdx);

        if (aPrimIdx == aData.y())
        {
          aMinPoint = aBox.CornerMin();
          aMaxPoint = aBox.CornerMax();
        }
        else
        {
          BVH::BoxMinMax<T, N>::CwiseMin (aMinPoint, aBox.CornerMin());
          BVH::BoxMinMax<T, N>::CwiseMax (aMaxPoint, aBox.CornerMax());
        }
      }
    }

    return 0;
  }
}

// =======================================================================
// function : EmitHierachy
// purpose  : Emits hierarchy from sorted Morton codes
// =======================================================================
template<class T, int N>
Standard_Integer BVH_LinearBuilder<T, N>::EmitHierachy (BVH_Tree<T, N>*                        theBVH,
                                                        const Standard_Integer                 theBit,
                                                        const Standard_Integer                 theShift,
                                                        std::vector<BVH_EncodedLink>::iterator theStart,
                                                        std::vector<BVH_EncodedLink>::iterator theFinal)
{
  if (theFinal - theStart > BVH_Builder<T, N>::myLeafNodeSize)
  {
    std::vector<BVH_EncodedLink>::iterator aPosition =
      (theBit >= 0) ? std::lower_bound (theStart, theFinal, BVH_EncodedLink(), BVH::BitComparator (theBit))
                    :                   theStart + ((theFinal - theStart) / 2);

    if (aPosition == theStart || aPosition == theFinal)
    {
      return EmitHierachy (theBVH, theBit - 1, theShift, theStart, theFinal);
    }

    // Build inner node
    const Standard_Integer aNode = theBVH->AddInnerNode (0, 0);

    const Standard_Integer aRghNode = theShift + static_cast<Standard_Integer> (aPosition - theStart);
    
    const Standard_Integer aLftChild = EmitHierachy (theBVH, theBit - 1, theShift, theStart, aPosition);
    const Standard_Integer aRghChild = EmitHierachy (theBVH, theBit - 1, aRghNode, aPosition, theFinal);

    theBVH->NodeInfoBuffer()[aNode].y() = aLftChild;
    theBVH->NodeInfoBuffer()[aNode].z() = aRghChild;
    
    return aNode;
  }
  else
  {
    // Build leaf node
    return theBVH->AddLeafNode (theShift, theShift + static_cast<Standard_Integer> (theFinal - theStart) - 1);
  }
}

#ifdef HAVE_TBB

namespace BVH
{
  //! TBB task for parallel radix sort.
  class RadixSortTask : public tbb::task
  {
    typedef std::vector<BVH_EncodedLink>::iterator LinkIterator;

  private:

    //! Start range element.
    LinkIterator myStart;
    
    //! Final range element.
    LinkIterator myFinal;

    //! Bit position for range partition.
    Standard_Integer myDigit;

  public:

    //! Creates new TBB radix sort task.
    RadixSortTask (LinkIterator theStart, LinkIterator theFinal, Standard_Integer theDigit)
    : myStart (theStart),
      myFinal (theFinal),
      myDigit (theDigit)
    {
      //
    }

    //! Executes the task.
    tbb::task* execute()
    {
      if (myDigit < 28)
      {
        BVH::RadixSorter::Perform (myStart, myFinal, myDigit);
      }
      else
      {
        LinkIterator anOffset = std::partition (myStart, myFinal, BitPredicate (myDigit));

        tbb::task_list aList;

        aList.push_back (*new ( allocate_child() )
          RadixSortTask (myStart, anOffset, myDigit - 1));

        aList.push_back (*new ( allocate_child() )
          RadixSortTask (anOffset, myFinal, myDigit - 1));

        set_ref_count (3); // count + 1
        spawn_and_wait_for_all (aList);
      }

      return NULL;
    }
  };

  //! TBB task for parallel bounds updating.
  template<class T, int N>
  class UpdateBoundTask: public tbb::task
  {
    //! Set of geometric objects.
    BVH_Set<T, N>* mySet;

    //! BVH tree built over the set.
    BVH_Tree<T, N>* myBVH;

    //! BVH node to update bounding box.
    Standard_Integer myNode;

    //! Level of the processed BVH node.
    Standard_Integer myLevel;

    //! Height of the processed BVH node.
    Standard_Integer* myHeight;

  public:

    //! Creates new TBB parallel bound update task.
    UpdateBoundTask (BVH_Set<T, N>*    theSet,
                     BVH_Tree<T, N>*   theBVH,
                     Standard_Integer  theNode,
                     Standard_Integer  theLevel,
                     Standard_Integer* theHeight)
    : mySet    (theSet),
      myBVH    (theBVH),
      myNode   (theNode),
      myLevel  (theLevel),
      myHeight (theHeight)
    {
      //
    }

    //! Executes the task.
    tbb::task* execute()
    {
      if (myBVH->IsOuter (myNode) || myLevel > 2)
      {
        *myHeight = BVH::UpdateBounds (mySet, myBVH, myNode);
      }
      else
      {
        Standard_Integer aLftHeight = 0;
        Standard_Integer aRghHeight = 0;

        tbb::task_list aList;

        const Standard_Integer aLftChild = myBVH->NodeInfoBuffer()[myNode].y();
        const Standard_Integer aRghChild = myBVH->NodeInfoBuffer()[myNode].z();

        Standard_Integer aCount = 1;

        if (!myBVH->IsOuter (aLftChild))
        {
          ++aCount;
          aList.push_back (*new ( allocate_child() )
            UpdateBoundTask (mySet, myBVH, aLftChild, myLevel + 1, &aLftHeight));
        }
        else
        {
          aLftHeight = BVH::UpdateBounds (mySet, myBVH, aLftChild);
        }

        if (!myBVH->IsOuter (aRghChild))
        {
          ++aCount;
          aList.push_back (*new( allocate_child() )
            UpdateBoundTask (mySet, myBVH, aRghChild, myLevel + 1, &aRghHeight));
        }
        else
        {
          aRghHeight = BVH::UpdateBounds (mySet, myBVH, aRghChild);
        }

        if (aCount > 1)
        {
          set_ref_count (aCount);
          spawn_and_wait_for_all (aList);
        }

        typename BVH_Box<T, N>::BVH_VecNt aLftMinPoint = myBVH->MinPointBuffer()[aLftChild];
        typename BVH_Box<T, N>::BVH_VecNt aLftMaxPoint = myBVH->MaxPointBuffer()[aLftChild];
        typename BVH_Box<T, N>::BVH_VecNt aRghMinPoint = myBVH->MinPointBuffer()[aRghChild];
        typename BVH_Box<T, N>::BVH_VecNt aRghMaxPoint = myBVH->MaxPointBuffer()[aRghChild];

        BVH::BoxMinMax<T, N>::CwiseMin (aLftMinPoint, aRghMinPoint);
        BVH::BoxMinMax<T, N>::CwiseMax (aLftMaxPoint, aRghMaxPoint);

        myBVH->MinPointBuffer()[myNode] = aLftMinPoint;
        myBVH->MaxPointBuffer()[myNode] = aLftMaxPoint;

        *myHeight = Max (aLftHeight, aRghHeight) + 1;
      }

      return NULL;
    }
  };
}

#endif

// =======================================================================
// function : Build
// purpose  :
// =======================================================================
template<class T, int N>
void BVH_LinearBuilder<T, N>::Build (BVH_Set<T, N>*       theSet,
                                     BVH_Tree<T, N>*      theBVH,
                                     const BVH_Box<T, N>& theBox)
{
  Standard_STATIC_ASSERT (N == 3 || N == 4);

  if (theBVH == NULL || theSet->Size() == 0)
  {
    return;
  }

  theBVH->Clear();

  const Standard_Integer aDimensionX = 1024;
  const Standard_Integer aDimensionY = 1024;
  const Standard_Integer aDimensionZ = 1024;

  const BVH_VecNt aSceneMin = theBox.CornerMin();
  const BVH_VecNt aSceneMax = theBox.CornerMax();

  const T aMinSize = static_cast<T> (BVH::THE_NODE_MIN_SIZE);

  const T aReverseSizeX = static_cast<T> (aDimensionX) / Max (aMinSize, aSceneMax.x() - aSceneMin.x());
  const T aReverseSizeY = static_cast<T> (aDimensionY) / Max (aMinSize, aSceneMax.y() - aSceneMin.y());
  const T aReverseSizeZ = static_cast<T> (aDimensionZ) / Max (aMinSize, aSceneMax.z() - aSceneMin.z());

  std::vector<BVH_EncodedLink> anEncodedLinks (theSet->Size(), BVH_EncodedLink());

  // Step 1 -- Assign Morton code to each primitive
  for (Standard_Integer aPrimIdx = 0; aPrimIdx < theSet->Size(); ++aPrimIdx)
  {
    const BVH_VecNt aCenter = theSet->Box (aPrimIdx).Center();

    Standard_Integer aVoxelX = BVH::IntFloor ((aCenter.x() - aSceneMin.x()) * aReverseSizeX);
    Standard_Integer aVoxelY = BVH::IntFloor ((aCenter.y() - aSceneMin.y()) * aReverseSizeY);
    Standard_Integer aVoxelZ = BVH::IntFloor ((aCenter.z() - aSceneMin.z()) * aReverseSizeZ);

    aVoxelX = Max (0, Min (aVoxelX, aDimensionX - 1));
    aVoxelY = Max (0, Min (aVoxelY, aDimensionY - 1));
    aVoxelZ = Max (0, Min (aVoxelZ, aDimensionZ - 1));

    aVoxelX = (aVoxelX | (aVoxelX << 16)) & 0x030000FF;
    aVoxelX = (aVoxelX | (aVoxelX <<  8)) & 0x0300F00F;
    aVoxelX = (aVoxelX | (aVoxelX <<  4)) & 0x030C30C3;
    aVoxelX = (aVoxelX | (aVoxelX <<  2)) & 0x09249249;

    aVoxelY = (aVoxelY | (aVoxelY << 16)) & 0x030000FF;
    aVoxelY = (aVoxelY | (aVoxelY <<  8)) & 0x0300F00F;
    aVoxelY = (aVoxelY | (aVoxelY <<  4)) & 0x030C30C3;
    aVoxelY = (aVoxelY | (aVoxelY <<  2)) & 0x09249249;

    aVoxelZ = (aVoxelZ | (aVoxelZ << 16)) & 0x030000FF;
    aVoxelZ = (aVoxelZ | (aVoxelZ <<  8)) & 0x0300F00F;
    aVoxelZ = (aVoxelZ | (aVoxelZ <<  4)) & 0x030C30C3;
    aVoxelZ = (aVoxelZ | (aVoxelZ <<  2)) & 0x09249249;

    anEncodedLinks[aPrimIdx] = BVH_EncodedLink (
      aVoxelX | (aVoxelY << 1) | (aVoxelZ << 2), aPrimIdx);
  }

  // Step 2 -- Sort primitives by their Morton codes using radix sort
#ifdef HAVE_TBB

  BVH::RadixSortTask& aSortTask = *new ( tbb::task::allocate_root() )
    BVH::RadixSortTask (anEncodedLinks.begin(), anEncodedLinks.end(), 29);

  tbb::task::spawn_root_and_wait (aSortTask);

#else

  BVH::RadixSorter::Perform (anEncodedLinks.begin(), anEncodedLinks.end());

#endif

  // Step 3 -- Emitting BVH hierarchy from sorted Morton codes
  EmitHierachy (theBVH, 29, 0, anEncodedLinks.begin(), anEncodedLinks.end());

  NCollection_Array1<Standard_Integer> aLinkMap (0, theSet->Size() - 1);

  for (Standard_Integer aLinkIdx = 0; aLinkIdx < theSet->Size(); ++aLinkIdx)
  {
    aLinkMap (anEncodedLinks[aLinkIdx].second) = aLinkIdx;
  }

  // Step 4 -- Rearranging primitive list according to Morton codes (in place)
  Standard_Integer aPrimIdx = 0;

  while (aPrimIdx < theSet->Size())
  {
    const Standard_Integer aSortIdx = aLinkMap (aPrimIdx);

    if (aPrimIdx != aSortIdx)
    {
      theSet->Swap (aPrimIdx, aSortIdx);

      std::swap (aLinkMap (aPrimIdx),
                 aLinkMap (aSortIdx));
    }
    else
    {
      ++aPrimIdx;
    }
  }

  // Step 5 -- Compute bounding boxes of BVH nodes
  theBVH->MinPointBuffer().resize (theBVH->NodeInfoBuffer().size());
  theBVH->MaxPointBuffer().resize (theBVH->NodeInfoBuffer().size());

  Standard_Integer aDepth = 0;

#ifdef HAVE_TBB

  // Note: Although TBB tasks are allocated using placement
  // new, we do not need to delete them explicitly
  BVH::UpdateBoundTask<T, N>& aRootTask = *new ( tbb::task::allocate_root() )
    BVH::UpdateBoundTask<T, N> (theSet, theBVH, 0, 0, &aDepth);

  tbb::task::spawn_root_and_wait (aRootTask);

#else

  aDepth = BVH::UpdateBounds (theSet, theBVH, 0);

#endif

  BVH_Builder<T, N>::UpdateDepth (theBVH, aDepth);
}