#pragma once

#include <cassert>
#include <cstdef> // For std::size_t.

#include <tbb/parallel_for.h>
#include <tbb/parallel_reduce.h>
#include <tbb/parallel_scan.h>
#include <tbb/parallel_sort.h>
#include <tbb/task_scheduler_init.h>
#include <tbb/tbb_thread.h>
#include <tbb/tick_count.h>

template <typename T>
struct NegateBody
{
  void operator()(T& x) const
  {
    x = -x;
  }
};

template <typename Vector>
struct ForBody
{
  typedef typename Vector::value_type T;

private:
  Vector& v;

public:
  ForBody(Vector& x)
      : v(x)
  {}

  void operator()(tbb::blocked_range<std::size_t> const& r) const
  {
    for (std::size_t i = r.begin(); i != r.end(); ++i)
    {
      v[i] = -v[i];
    }
  }
};

template <typename Vector>
struct ReduceBody
{
  typedef typename Vector::value_type T;

private:
  Vector& v;

public:
  T sum;

  ReduceBody(Vector& x)
      : v(x)
      , sum(0)
  {}

  ReduceBody(ReduceBody& x, tbb::split)
      : v(x.v)
      , sum(0)
  {}

  void operator()(tbb::blocked_range<std::size_t> const& r)
  {
    for (std::size_t i = r.begin(); i != r.end(); ++i)
    {
      sum += v[i];
    }
  }

  void join(ReduceBody const& x)
  {
    sum += x.sum;
  }
};

template <typename Vector>
struct ScanBody
{
  typedef typename Vector::value_type T;

private:
  Vector& v;

public:
  T sum;

  ScanBody(Vector& x)
      : sum(0)
      , v(x)
  {}

  ScanBody(ScanBody& x, tbb::split)
      : v(x.v)
      , sum(0)
  {}

  template <typename Tag>
  void operator()(tbb::blocked_range<std::size_t> const& r, Tag)
  {
    T temp = sum;
    for (std::size_t i = r.begin(); i < r.end(); ++i)
    {
      temp = temp + x[i];
      if (Tag::is_final_scan())
      {
        x[i] = temp;
      }
    }
    sum = temp;
  }

  void assign(ScanBody const& x)
  {
    sum = x.sum;
  }

  T get_sum() const
  {
    return sum;
  }

  void reverse_join(ScanBody const& x)
  {
    sum = x.sum + sum;
  }
};

template <typename Vector>
struct CopyBody
{
  typedef typename Vector::value_type T;

private:
  Vector& v;
  Vector& u;

public:
  CopyBody(Vector& x, Vector& y)
      : v(x)
      , u(y)
  {}

  void operator()(tbb::blocked_range<size_t> const& r) const
  {
    for (std::size_t i = r.begin(); i != r.end(); ++i)
    {
      v[i] = u[i];
    }
  }
};

template <typename Vector>
typename Vector::value_type tbb_reduce(Vector& v)
{
  ReduceBody<Vector> body(v);
  tbb::parallel_reduce(tbb::blocked_range<size_t>(0, v.size()), body);
  return body.sum;
}

template <typename Vector>
void tbb_sort(Vector& v)
{
  tbb::parallel_sort(v.begin(), v.end());
}

template <typename Vector>
void tbb_transform(Vector& v)
{
  ForBody<Vector> body(v);
  tbb::parallel_for(tbb::blocked_range<size_t>(0, v.size()), body);
}

template <typename Vector>
void tbb_scan(Vector& v)
{
  ScanBody<Vector> body(v);
  tbb::parallel_scan(tbb::blocked_range<size_t>(0, v.size()), body);
}

template <typename Vector>
void tbb_copy(Vector& v, Vector& u)
{
  CopyBody<Vector> body(v, u);
  tbb::parallel_for(tbb::blocked_range<size_t>(0, v.size()), body);
}

void test_tbb()
{
  std::size_t elements = 1 << 20;

  std::vector<int> A(elements);
  std::vector<int> B(elements);
  std::vector<int> C(elements);
  std::vector<int> D(elements);

  randomize(A);
  randomize(B);
  assert(std::accumulate(A.begin(), A.end(), 0) == tbb_reduce(A));

  randomize(A);
  randomize(B);
  std::transform(A.begin(), A.end(), A.begin(), thrust::negate<int>());
  tbb_transform(B);
  assert(A == B);

  randomize(A);
  randomize(B);
  std::partial_sum(A.begin(), A.end(), A.begin());
  tbb_scan(B);
  assert(A == B);

  randomize(A);
  randomize(B);
  std::sort(A.begin(), A.end());
  tbb_sort(B);
  assert(A == B);

  randomize(A);
  randomize(B);
  randomize(C);
  randomize(D);
  std::copy(A.begin(), A.end(), C.begin());
  tbb_copy(B, D);
  assert(A == B);
  assert(C == D);
}
