// Copyright (c) 2018 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 <vector>
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/detail/safe_ref.h"
#include "paddle/fluid/platform/device_context.h"
#include "paddle/fluid/platform/for_range.h"
#ifdef PADDLE_WITH_CUDA
#include <thrust/random.h>
#endif

namespace paddle {
namespace operators {

template <typename DeviceContext>
struct Random;

template <>
struct Random<platform::CPUDeviceContext> {
  using Engine = std::minstd_rand;

  template <typename T>
  using UniformIntDist = std::uniform_int_distribution<T>;
};

#ifdef PADDLE_WITH_CUDA
template <>
struct Random<platform::CUDADeviceContext> {
  using Engine = thrust::minstd_rand;

  template <typename T>
  using UniformIntDist = thrust::uniform_int_distribution<T>;
};
#endif

template <typename T>
HOSTDEVICE inline void StridedMemcpy(const T* x, const size_t* x_dims, T* out,
                                     const size_t* out_dims, int i, int rank,
                                     size_t prod_x_remain,
                                     size_t prod_out_remain,
                                     const size_t* offsets) {
  size_t x_dim_i = x_dims[i];
  size_t out_dim_i = out_dims[i];
  size_t x_stride = prod_x_remain / x_dim_i;
  size_t out_stride = prod_out_remain / out_dim_i;
  size_t offset_i = offsets[i];

  if (i == rank - 1) {
    PADDLE_ASSERT(x_stride == 1 && out_stride == 1);
    x += offset_i;
    for (size_t j = 0; j < out_dim_i; ++j) {
      *out++ = *x++;
    }
  } else {
    x += offset_i * x_stride;
    for (size_t j = 0; j < out_dim_i; ++j) {
      StridedMemcpy<T>(x, x_dims, out, out_dims, i + 1, rank, x_stride,
                       out_stride, offsets);
      x += x_stride;
      out += out_stride;
    }
  }
}

template <typename DeviceContext, typename T>
struct RandomCropFunctor {
  const T* x_;
  T* out_;
  size_t x_dims_[9];
  size_t out_dims_[9];
  int num_batchsize_dims_;
  int rank_;
  int64_t seed_;

  size_t prod_batchsize_dims_;
  size_t prod_x_ins_dims_;
  size_t prod_out_ins_dims_;

  RandomCropFunctor(const T* x, T* out, const framework::DDim& x_dims,
                    const framework::DDim& out_dims, int num_batchsize_dims,
                    int64_t seed)
      : x_(x),
        out_(out),
        num_batchsize_dims_(num_batchsize_dims),
        rank_(x_dims.size()),
        seed_(seed) {
    PADDLE_ENFORCE_EQ(x_dims.size(), out_dims.size());
    PADDLE_ENFORCE_GT(rank_, num_batchsize_dims_);
    prod_batchsize_dims_ = 1;
    prod_x_ins_dims_ = 1;
    prod_out_ins_dims_ = 1;
    for (size_t i = 0; i < static_cast<size_t>(rank_); ++i) {
      size_t x_dim_i = x_dims[i];
      size_t out_dim_i = out_dims[i];
      x_dims_[i] = x_dim_i;
      out_dims_[i] = out_dim_i;
      if (i < static_cast<size_t>(num_batchsize_dims_)) {
        PADDLE_ENFORCE_EQ(x_dim_i, out_dim_i);
        prod_batchsize_dims_ *= x_dim_i;
      } else {
        prod_x_ins_dims_ *= x_dim_i;
        prod_out_ins_dims_ *= out_dim_i;
      }
    }
  }

  HOSTDEVICE void operator()(size_t ins_idx) {
    typename Random<DeviceContext>::Engine engine(seed_);
    engine.discard(ins_idx * (rank_ - num_batchsize_dims_));
    size_t offsets[9] = {};
    for (int i = num_batchsize_dims_; i < rank_; ++i) {
      typename Random<DeviceContext>::template UniformIntDist<size_t> dist(
          0, x_dims_[i] - out_dims_[i]);
      offsets[i - num_batchsize_dims_] = dist(engine);
    }

    const T* x = x_ + ins_idx * prod_x_ins_dims_;
    T* out = out_ + ins_idx * prod_out_ins_dims_;

    StridedMemcpy<T>(x, x_dims_ + num_batchsize_dims_, out,
                     out_dims_ + num_batchsize_dims_, 0,
                     rank_ - num_batchsize_dims_, prod_x_ins_dims_,
                     prod_out_ins_dims_, offsets);
  }
};

template <typename DeviceContext, typename T>
class RandomCropKernel : public framework::OpKernel<T> {
 public:
  virtual void Compute(const framework::ExecutionContext& ctx) const {
    int64_t seed = 0;
    auto& seed_tensor = detail::Ref(ctx.Input<framework::LoDTensor>("Seed"));
    if (seed_tensor.IsInitialized()) {
      if (platform::is_cpu_place(seed_tensor.place())) {
        seed = *seed_tensor.data<int64_t>();
      } else {
        LOG(WARNING) << "It is slow to place seed in GPU memory. Please verify "
                        "your program";
        framework::LoDTensor cpu_seed;
        framework::TensorCopySync(seed_tensor, platform::CPUPlace(), &cpu_seed);
        seed = *cpu_seed.data<int64_t>();
      }
    } else {
      VLOG(5) << "WARNING: The input 'Seed' is not initialized, use attribute "
                 "'startup_seed' instead.";
      seed = ctx.Attr<int>("startup_seed");
    }
    auto shape = ctx.Attr<std::vector<int>>("shape");
    auto& x = detail::Ref(ctx.Input<framework::LoDTensor>("X"));
    auto& out = detail::Ref(ctx.Output<framework::LoDTensor>("Out"));

    int num_batchsize_dims = x.dims().size() - shape.size();
    RandomCropFunctor<DeviceContext, T> functor(
        x.data<T>(), out.mutable_data<T>(ctx.GetPlace()), x.dims(), out.dims(),
        num_batchsize_dims, seed);
    platform::ForRange<DeviceContext> for_range(
        ctx.template device_context<DeviceContext>(),
        functor.prod_batchsize_dims_);

    for_range(functor);

    Random<platform::CPUDeviceContext>::Engine engine(seed);
    engine.discard(functor.prod_batchsize_dims_ *
                   (functor.rank_ - functor.num_batchsize_dims_));
    *ctx.Output<framework::LoDTensor>("SeedOut")->mutable_data<int64_t>(
        framework::make_ddim({1}), platform::CPUPlace()) = engine();
  }
};

// TODO(fengjiayi): Backward of random crop op

}  // namespace operators
}  // namespace paddle