// ======================================================================== //
// Copyright 2009-2015 Intel Corporation                                    //
//                                                                          //
// Licensed under the Apache License, Version 2.0 (the "License");          //
// you may not use this file except in compliance with the License.         //
// You may obtain a copy of the License at                                  //
//                                                                          //
//     http://www.apache.org/licenses/LICENSE-2.0                           //
//                                                                          //
// Unless required by applicable law or agreed to in writing, software      //
// distributed under the License is distributed on an "AS IS" BASIS,        //
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. //
// See the License for the specific language governing permissions and      //
// limitations under the License.                                           //
// ======================================================================== //

#pragma once

#include "../common/default.h"
#include "parallel_for.h"

namespace embree
{
  /*! Implementation of parallel prefix operations (e.g. parallel
   *  prefix sums). The prefix operations start with the identity
   *  (e.g. zero for prefix sums). */
  template<typename SrcArray, typename DstArray, typename Ty, typename Op>
  class ParallelPrefixOp
  {
#if defined(__MIC__)
    static const size_t MAX_TASKS = MAX_THREADS;
    static const size_t SINGLE_THREAD_THRESHOLD = 50000;
#else
    static const size_t MAX_TASKS = MAX_THREADS;
    static const size_t SINGLE_THREAD_THRESHOLD = 3000000;
#endif

  public:
    ParallelPrefixOp () {}

    class Task
    {
    public:

      Task (ParallelPrefixOp* parent, const SrcArray& src, DstArray& dst, const size_t N, const Op op, const Ty id)
	: parent(parent), src(src), dst(dst), N(N), op(op), id(id)
      {
	/* perform single threaded prefix operation for small N */
	if (N < SINGLE_THREAD_THRESHOLD) {
          size_t sum=0;
	  for (size_t i=0; i<N; sum+=src[i++]) dst[i] = sum;
          parent->value = sum;
	}

	/* perform parallel prefix operation for large N */
	else 
	{
	  const size_t threadCount = min(MAX_TASKS, TaskSchedulerTBB::threadCount());
          
	  /* first calculate range for each block */
	  parallel_for(threadCount, [&] (const size_t threadIndex) 
	  {
	    const size_t start = (threadIndex+0)*N/threadCount;
	    const size_t end   = (threadIndex+1)*N/threadCount;
	    
	    Ty sum = id;
	    for (size_t i=start; i<end; i++)
	      sum = op(sum,src[i]);
	    
	    parent->state[threadIndex] = sum;
	  });

	  /* now calculate prefix_op for each block */
	  parallel_for(threadCount, [&] (const size_t threadIndex) 
	  {
	    const size_t start = (threadIndex+0)*N/threadCount;
	    const size_t end   = (threadIndex+1)*N/threadCount;
	    
	    /* calculate start sum for block */
	    Ty count = id;
	    for (size_t i=0; i<threadIndex; i++)
	      count += parent->state[i];
	    
	    /* calculate prefix sums for the block */
	    for (size_t i=start; i<end; i++) 
	    {
	      const Ty v = src[i];
	      dst[i] = count;
	      count = op(count,v);
	    }
	    
	    if (threadIndex == threadCount-1) 
	      parent->value = count;
	  });
	}
      }
      
    private:
      ParallelPrefixOp* const parent;
      const SrcArray& src;               //!< source array
      DstArray& dst;                     //!< destination array
      const size_t N;                    //!< number of elements in the arrays
      const Op op;                       //!< operation to use for prefix sum
      const Ty id;                       //!< identity of the operation
    };

    void operator() (const SrcArray& src, DstArray& dst, const size_t N, const Op op, const Ty id) {
      Task(this,src,dst,N,op,id);
    }

  public:
    Ty state[MAX_TASKS];
    Ty value;
  };

  /*! parallel calculation of prefix sums */
  template<typename SrcArray, typename DstArray, typename Add>
    __forceinline typename SrcArray::value_type parallel_prefix_sum(const SrcArray& src, DstArray& dst, size_t N, const Add& add) 
  {
    typedef typename SrcArray::value_type Value;
    ParallelPrefixOp<SrcArray,DstArray,Value,Add> op; op(src,dst,N,add,0); return op.value;
  }
}
