// Copyright (c) 2024 PaddlePaddle Authors. 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.

#pragma once
#include <algorithm>
#include <utility>
#include <vector>

#include "paddle/phi/backends/context_pool.h"
#if defined(__NVCC__) || defined(__HIPCC__)
#include <thrust/random.h>

#include "paddle/phi/core/generator.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/funcs/distribution_helper.h"
#include "paddle/phi/kernels/funcs/index_impl.cu.h"
#endif

#include "glog/logging.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/tensor_utils.h"

namespace phi {
namespace funcs {

template <typename T>
inline void UniformRealDistribution(T* data,
                                    const int64_t& size,
                                    const float& min,
                                    const float& max,
                                    const unsigned int seed) {
  VLOG(4) << "[CPU] UniformRandomKernel<T>";
  std::uniform_real_distribution<T> dist(static_cast<T>(min),
                                         static_cast<T>(max));
  auto engine = phi::GetCPURandomEngine(seed);

  for (int64_t i = 0; i < size; ++i) {
    data[i] = dist(*engine);
  }
}

template <>
inline void UniformRealDistribution(phi::bfloat16* data,
                                    const int64_t& size,
                                    const float& min,
                                    const float& max,
                                    const unsigned int seed) {
  VLOG(4) << "[CPU] UniformRandomKernel<bfloat16>";
  std::uniform_real_distribution<float> dist(min, max);
  auto engine = phi::GetCPURandomEngine(seed);

  for (int64_t i = 0; i < size; ++i) {
    data[i] = static_cast<phi::bfloat16>(dist(*engine));
  }
}

inline std::vector<int64_t> GetNewDataFromShapeTensor(
    const DenseTensor* new_data_tensor) {
  DenseTensor cpu_starts_tensor;
  auto* dev_ctx = DeviceContextPool::Instance().Get(cpu_starts_tensor.place());
  if (new_data_tensor->dtype() == DataType::INT64) {
    auto* new_data = new_data_tensor->data<int64_t>();
    if (new_data_tensor->place().GetType() == AllocationType::GPU) {
      phi::Copy(
          *dev_ctx, *new_data_tensor, CPUPlace(), true, &cpu_starts_tensor);
      new_data = cpu_starts_tensor.data<int64_t>();
    }
    std::vector<int64_t> vec_new_data(new_data,
                                      new_data + new_data_tensor->numel());
    return vec_new_data;
  } else if (new_data_tensor->dtype() == DataType::INT32) {
    auto* new_data = new_data_tensor->data<int32_t>();
    std::vector<int64_t> vec_new_data;
    if (new_data_tensor->place().GetType() == AllocationType::GPU) {
      phi::Copy(
          *dev_ctx, *new_data_tensor, CPUPlace(), true, &cpu_starts_tensor);
      new_data = cpu_starts_tensor.data<int32_t>();
    }
    for (int64_t i = 0; i < new_data_tensor->numel(); ++i) {
      vec_new_data.push_back(static_cast<int64_t>(*(new_data + i)));
    }
    return vec_new_data;
  } else {
    PADDLE_THROW(common::errors::InvalidArgument(
        "Expected dtype of ShapeTensor must be int32, int64. But got "
        "unsupported dtype: %s.",
        new_data_tensor->dtype()));
  }
}

inline std::vector<int64_t> GetNewDataFromShapeTensorList(
    const std::vector<const DenseTensor*>& list_new_shape_tensor) {
  DenseTensor temp;
  auto* dev_ctx = DeviceContextPool::Instance().Get(temp.place());
  std::vector<int64_t> vec_new_shape;
  vec_new_shape.reserve(list_new_shape_tensor.size());
  for (size_t i = 0; i < list_new_shape_tensor.size(); ++i) {
    auto tensor = list_new_shape_tensor[i];
    PADDLE_ENFORCE_EQ(
        tensor->dims(),
        make_ddim({1}),
        common::errors::InvalidArgument(
            "Shape of dim tensor in uniform_random_op should be [1]"
            "But received tensor's dim=%s.",
            tensor->dims()));

    if (tensor->dtype() == DataType::INT32) {
      if (tensor->place().GetType() == AllocationType::GPU) {
        phi::Copy(*dev_ctx, *tensor, CPUPlace(), true, &temp);
        vec_new_shape.push_back(static_cast<int64_t>(*temp.data<int32_t>()));
      } else {
        vec_new_shape.push_back(static_cast<int64_t>(*tensor->data<int32_t>()));
      }
    } else if (tensor->dtype() == DataType::INT64) {
      if (tensor->place().GetType() == AllocationType::GPU) {
        DenseTensor temp;
        phi::Copy(*dev_ctx, *tensor, CPUPlace(), true, &temp);
        vec_new_shape.push_back(*temp.data<int64_t>());
      } else {
        vec_new_shape.push_back(*tensor->data<int64_t>());
      }
    } else {
      PADDLE_THROW(common::errors::InvalidArgument(
          "Expected dtype of ShapeTensorList of %d-th must be int32, int64. "
          "But got "
          "unsupported dtype: %s.",
          i,
          DataTypeToString(tensor->dtype())));
    }
  }

  return vec_new_shape;
}

#if defined(__NVCC__) || defined(__HIPCC__)
template <typename T>
struct UniformGenerator {
  T min_, max_;
  unsigned int seed_;
  T diag_val_;
  unsigned int diag_num_;
  unsigned int diag_step_;
  __host__ __device__ UniformGenerator(
      T min, T max, int seed, int diag_num, int diag_step, T diag_val)
      : min_(min),
        max_(max),
        seed_(seed),
        diag_num_(diag_num),
        diag_step_(diag_step),
        diag_val_(diag_val) {}

  __host__ __device__ T operator()(const unsigned int n) const {
    thrust::minstd_rand rng;
    rng.seed(seed_);
    thrust::uniform_real_distribution<T> dist(min_, max_);
    rng.discard(n);
    T out = dist(rng);
    unsigned int remainder = n % (diag_step_ + 1);
    if (remainder == 0 && diag_num_ > n / (diag_step_ + 1)) {
      out = diag_val_;
    }
    return out;
  }
};

template <typename T>
void UniformRandom(const GPUContext& dev_ctx,
                   DenseTensor* tensor,
                   int attr_seed,
                   float attr_min,
                   float attr_max,
                   int attr_diag_num,
                   int attr_diag_step,
                   float attr_diag_val) {
  int64_t size = tensor->numel();
  T* data = dev_ctx.Alloc<T>(tensor);
  if (size <= 0) return;
  unsigned int seed = static_cast<unsigned int>(attr_seed);

  T min = static_cast<T>(attr_min);
  T max = static_cast<T>(attr_max);
  unsigned int diag_num = static_cast<unsigned int>(attr_diag_num);
  unsigned int diag_step = static_cast<unsigned int>(attr_diag_step);
  T diag_val = static_cast<T>(attr_diag_val);

  if (seed == 0) {
    // Use global Generator seed
    using MT = typename MPTypeTrait<T>::Type;
    funcs::uniform_distribution<MT> dist;
    funcs::uniform_real_transform<MT> trans(min, max);
    funcs::distribution_and_transform<T>(dev_ctx, tensor, dist, trans);
  } else {
    // Use OP seed
    auto func =
        UniformGenerator<T>(min, max, seed, diag_num, diag_step, diag_val);
    phi::IndexKernel<T, UniformGenerator<T>>(dev_ctx, tensor, func);
  }
}
#endif
}  // namespace funcs
}  // namespace phi