/*
 *  Copyright 2008-2013 NVIDIA Corporation
 *  Modifications Copyright© 2019-2025 Advanced Micro Devices, Inc. All rights reserved.
 *
 *  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.
 */

#include <thrust/functional.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/permutation_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/sort.h>

#include "test_real_assertions.hpp"
#include "test_param_fixtures.hpp"
#include "test_utils.hpp"

TESTS_DEFINE(SortPermutationIteratorsTests, FullTestsParams);

template <typename Iterator>
class strided_range
{
public:
  using difference_type = typename thrust::iterator_difference<Iterator>::type;

  struct stride_functor
  {
    difference_type stride;

    stride_functor(difference_type stride)
        : stride(stride)
    {}

    __host__ __device__ difference_type operator()(const difference_type& i) const
    {
      return stride * i;
    }
  };

  using CountingIterator    = typename thrust::counting_iterator<difference_type>;
  using TransformIterator   = typename thrust::transform_iterator<stride_functor, CountingIterator>;
  using PermutationIterator = typename thrust::permutation_iterator<Iterator, TransformIterator>;

  // type of the strided_range iterator
  using iterator = PermutationIterator;

  // construct strided_range for the range [first,last)
  strided_range(Iterator first, Iterator last, difference_type stride)
      : first(first)
      , last(last)
      , stride(stride)
  {}

  iterator begin(void) const
  {
    return PermutationIterator(first, TransformIterator(CountingIterator(0), stride_functor(stride)));
  }

  iterator end(void) const
  {
    return begin() + ((last - first) + (stride - 1)) / stride;
  }

protected:
  Iterator first;
  Iterator last;
  difference_type stride;
};

TYPED_TEST(SortPermutationIteratorsTests, TestSortPermutationIterator)
{
  using Vector   = typename TestFixture::input_type;
  using Policy   = typename TestFixture::execution_policy;
  using T        = typename Vector::value_type;
  using Iterator = typename Vector::iterator;

  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());

  Vector A(10);
  A[0] = T(2);
  A[1] = T(9);
  A[2] = T(0);
  A[3] = T(1);
  A[4] = T(5);
  A[5] = T(3);
  A[6] = T(8);
  A[7] = T(6);
  A[8] = T(7);
  A[9] = T(4);

  strided_range<Iterator> S(A.begin(), A.end(), 2);

  thrust::sort(Policy{}, S.begin(), S.end());

  ASSERT_EQ(A[0], T(0));
  ASSERT_EQ(A[1], T(9));
  ASSERT_EQ(A[2], T(2));
  ASSERT_EQ(A[3], T(1));
  ASSERT_EQ(A[4], T(5));
  ASSERT_EQ(A[5], T(3));
  ASSERT_EQ(A[6], T(7));
  ASSERT_EQ(A[7], T(6));
  ASSERT_EQ(A[8], T(8));
  ASSERT_EQ(A[9], T(4));
}

TYPED_TEST(SortPermutationIteratorsTests, TestStableSortPermutationIterator)
{
  using Vector   = typename TestFixture::input_type;
  using Policy   = typename TestFixture::execution_policy;
  using T        = typename Vector::value_type;
  using Iterator = typename Vector::iterator;

  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());

  Vector A(10);
  A[0] = T(2);
  A[1] = T(9);
  A[2] = T(0);
  A[3] = T(1);
  A[4] = T(5);
  A[5] = T(3);
  A[6] = T(8);
  A[7] = T(6);
  A[8] = T(7);
  A[9] = T(4);

  strided_range<Iterator> S(A.begin(), A.end(), 2);

  thrust::stable_sort(Policy{}, S.begin(), S.end());

  ASSERT_EQ(A[0], T(0));
  ASSERT_EQ(A[1], T(9));
  ASSERT_EQ(A[2], T(2));
  ASSERT_EQ(A[3], T(1));
  ASSERT_EQ(A[4], T(5));
  ASSERT_EQ(A[5], T(3));
  ASSERT_EQ(A[6], T(7));
  ASSERT_EQ(A[7], T(6));
  ASSERT_EQ(A[8], T(8));
  ASSERT_EQ(A[9], T(4));
}

TYPED_TEST(SortPermutationIteratorsTests, TestSortByKeyPermutationIterator)
{
  using Vector    = typename TestFixture::input_type;
  using Policy    = typename TestFixture::execution_policy;
  using ValueType = typename Vector::value_type;
  using Iterator  = typename Vector::iterator;

  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());

  Vector A(10), B(10);
  A[0] = ValueType(2);
  B[0] = ValueType(0);
  A[1] = ValueType(9);
  B[1] = ValueType(1);
  A[2] = ValueType(0);
  B[2] = ValueType(2);
  A[3] = ValueType(1);
  B[3] = ValueType(3);
  A[4] = ValueType(5);
  B[4] = ValueType(4);
  A[5] = ValueType(3);
  B[5] = ValueType(5);
  A[6] = ValueType(8);
  B[6] = ValueType(6);
  A[7] = ValueType(6);
  B[7] = ValueType(7);
  A[8] = ValueType(7);
  B[8] = ValueType(8);
  A[9] = ValueType(4);
  B[9] = ValueType(9);

  strided_range<Iterator> S(A.begin(), A.end(), 2);
  strided_range<Iterator> T(B.begin(), B.end(), 2);

  thrust::sort_by_key(Policy{}, S.begin(), S.end(), T.begin());

  ASSERT_EQ(A[0], ValueType(0));
  ASSERT_EQ(A[1], ValueType(9));
  ASSERT_EQ(A[2], ValueType(2));
  ASSERT_EQ(A[3], ValueType(1));
  ASSERT_EQ(A[4], ValueType(5));
  ASSERT_EQ(A[5], ValueType(3));
  ASSERT_EQ(A[6], ValueType(7));
  ASSERT_EQ(A[7], ValueType(6));
  ASSERT_EQ(A[8], ValueType(8));
  ASSERT_EQ(A[9], ValueType(4));

  ASSERT_EQ(B[0], ValueType(2));
  ASSERT_EQ(B[1], ValueType(1));
  ASSERT_EQ(B[2], ValueType(0));
  ASSERT_EQ(B[3], ValueType(3));
  ASSERT_EQ(B[4], ValueType(4));
  ASSERT_EQ(B[5], ValueType(5));
  ASSERT_EQ(B[6], ValueType(8));
  ASSERT_EQ(B[7], ValueType(7));
  ASSERT_EQ(B[8], ValueType(6));
  ASSERT_EQ(B[9], ValueType(9));
}

TYPED_TEST(SortPermutationIteratorsTests, TestStableSortByKeyPermutationIterator)
{
  using Vector    = typename TestFixture::input_type;
  using Policy    = typename TestFixture::execution_policy;
  using ValueType = typename Vector::value_type;
  using Iterator  = typename Vector::iterator;

  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());

  Vector A(10), B(10);
  A[0] = ValueType(2);
  B[0] = ValueType(0);
  A[1] = ValueType(9);
  B[1] = ValueType(1);
  A[2] = ValueType(0);
  B[2] = ValueType(2);
  A[3] = ValueType(1);
  B[3] = ValueType(3);
  A[4] = ValueType(5);
  B[4] = ValueType(4);
  A[5] = ValueType(3);
  B[5] = ValueType(5);
  A[6] = ValueType(8);
  B[6] = ValueType(6);
  A[7] = ValueType(6);
  B[7] = ValueType(7);
  A[8] = ValueType(7);
  B[8] = ValueType(8);
  A[9] = ValueType(4);
  B[9] = ValueType(9);

  strided_range<Iterator> S(A.begin(), A.end(), 2);
  strided_range<Iterator> T(B.begin(), B.end(), 2);

  thrust::stable_sort_by_key(Policy{}, S.begin(), S.end(), T.begin());

  ASSERT_EQ(A[0], ValueType(0));
  ASSERT_EQ(A[1], ValueType(9));
  ASSERT_EQ(A[2], ValueType(2));
  ASSERT_EQ(A[3], ValueType(1));
  ASSERT_EQ(A[4], ValueType(5));
  ASSERT_EQ(A[5], ValueType(3));
  ASSERT_EQ(A[6], ValueType(7));
  ASSERT_EQ(A[7], ValueType(6));
  ASSERT_EQ(A[8], ValueType(8));
  ASSERT_EQ(A[9], ValueType(4));

  ASSERT_EQ(B[0], ValueType(2));
  ASSERT_EQ(B[1], ValueType(1));
  ASSERT_EQ(B[2], ValueType(0));
  ASSERT_EQ(B[3], ValueType(3));
  ASSERT_EQ(B[4], ValueType(4));
  ASSERT_EQ(B[5], ValueType(5));
  ASSERT_EQ(B[6], ValueType(8));
  ASSERT_EQ(B[7], ValueType(7));
  ASSERT_EQ(B[8], ValueType(6));
  ASSERT_EQ(B[9], ValueType(9));
}
