// 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.

#include <xxhash.h>
#include <algorithm>
#include <cmath>
#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/math/bloomfilter.h"
#include "paddle/phi/kernels/funcs/search_compute.h"

namespace phi {

#ifndef _WIN32
bool should_use_term(math::bloomfilter* _filter,
                     math::bloomfilter* _black_filter,
                     const float* word_repr,
                     int len) {
  return (!_filter || 1 == math::bloomfilter_get(
                               _filter, word_repr, len * sizeof(float))) &&
         (!_black_filter ||
          0 == math::bloomfilter_get(
                   _black_filter, word_repr, len * sizeof(float)));
}

template <typename T>
void hash_embedding_ff(const float* hash_id,
                       int len,
                       T* top_pos,
                       const T* weights,
                       int _num_emb,
                       int _rand_len,
                       int _space_len) {
  unsigned int pos1 = XXH32(hash_id, len * sizeof(float), 0) % _space_len;
  unsigned int pos2 =
      XXH32(hash_id, len * sizeof(float), _rand_len) % _space_len;

  for (int j = 0; j != _num_emb; j += _rand_len) {
    if (j + _rand_len < _num_emb) {
      __builtin_prefetch(weights + pos2);
      __builtin_prefetch(top_pos + j + _rand_len);
    }

    unsigned int pos3 =
        XXH32(hash_id, len * sizeof(float), j + 2 * _rand_len) % _space_len;
    memcpy(top_pos + j, const_cast<T*>(weights + pos1), _rand_len * sizeof(T));
    pos1 = pos2;
    pos2 = pos3;
  }
}

template <typename T, typename Context>
void CPUPyramidHashOPKernel(const Context& dev_ctx,
                            const DenseTensor& x,
                            const DenseTensor& w,
                            const DenseTensor& white_list,
                            const DenseTensor& black_list,
                            int num_emb,
                            int space_len,
                            int pyramid_layer,
                            int rand_len,
                            float drop_out_percent,
                            int is_training,
                            bool use_filter,
                            int white_list_len,
                            int black_list_len,
                            int seed,
                            float lr,
                            const std::string& distribute_update_vars,
                            DenseTensor* out,
                            DenseTensor* drop_pos,
                            DenseTensor* x_temp_out) {
  auto* bottom = &x;
  auto* _blobs_0 = &w;
  auto* _blobs_1 = &white_list;
  auto* _blobs_2 = &black_list;
  auto* top = out;

  int _num_emb = num_emb;
  int _pyramid_layer = pyramid_layer;
  int _is_training = is_training;
  unsigned int _seed = (unsigned int)seed;
  int _rand_len = rand_len;
  int _space_len = space_len;
  float _drop_out_percent = drop_out_percent;

  const auto& offset = bottom->lod()[0];
  const auto* bottom_data_ori = bottom->data<int32_t>();
  auto* buff = x_temp_out;
  buff->Resize({bottom->dims()[0], bottom->dims()[1]});
  float* bottom_data = dev_ctx.template Alloc<float>(buff);
  for (int i = 0; i < bottom->dims()[0]; i++) {
    bottom_data[i] = bottom_data_ori[i];  // NOLINT
  }

  const auto* weights = _blobs_0->data<T>();

  std::vector<size_t> top_offset;
  top_offset.resize(offset.size());
  top_offset[0] = 0;

  math::bloomfilter* _filter = nullptr;
  math::bloomfilter* _black_filter = nullptr;
  if (use_filter) {
    if (white_list_len != 0) {
      _filter = (math::bloomfilter*)_blobs_1->data<float>();
      PADDLE_ENFORCE_EQ(
          math::bloomfilter_check(_filter),
          1,
          common::errors::PreconditionNotMet(
              "The white filter is not loaded successfully, please make sure "
              "'white_list_len': %d is valid for Input(WhiteList).",
              white_list_len));
    }
    if (black_list_len != 0) {
      _black_filter = (math::bloomfilter*)_blobs_2->data<float>();
      PADDLE_ENFORCE_EQ(
          math::bloomfilter_check(_black_filter),
          1,
          common::errors::PreconditionNotMet(
              "The black filter is not loaded successfully, please make sure "
              "'black_list_len': %d is valid for Input(BlackList).",
              black_list_len));
    }
  }

  drop_pos->Resize(
      make_ddim({bottom->dims()[0] * bottom->dims()[1] * _pyramid_layer, 1}));
  std::vector<size_t> drop_pos_offset;
  drop_pos_offset.resize(offset.size());
  drop_pos_offset[0] = 0;
  int* iter = dev_ctx.template Alloc<int>(drop_pos);
  int* iter_end = iter;

  for (size_t i = 0; i < top_offset.size() - 1; ++i) {
    int w = static_cast<int>(offset[i + 1] - offset[i]);
    int nsentense_with_pyramid = 0;
    if (w < 2) {
      nsentense_with_pyramid = 0;
    } else {
      for (int ilayer = 1; ilayer < _pyramid_layer && ilayer < w; ++ilayer) {
        for (int l = 0; l < w - ilayer; ++l) {
          if (should_use_term(_filter,
                              _black_filter,
                              (const float*)(bottom_data + offset[i] + l),
                              ilayer + 1)) {
            if (_is_training != 0) {
              unsigned int rand_val = rand_r(&_seed);
              double rate = static_cast<double>(rand_val) / (RAND_MAX);
              *(iter_end++) = (rate < _drop_out_percent ? 0 : 1);
            } else {
              *(iter_end++) = 1;
            }
          } else {
            *(iter_end++) = 0;
          }
        }
      }
      nsentense_with_pyramid = static_cast<int>(std::count(iter, iter_end, 1));
      iter = iter_end;
    }
    drop_pos_offset[i + 1] = drop_pos_offset[i] + nsentense_with_pyramid;
    top_offset[i + 1] =
        top_offset[i] +
        (nsentense_with_pyramid == 0 ? 1 : nsentense_with_pyramid);
  }

  int top_l = static_cast<int>(top_offset[top_offset.size() - 1]);

  LegacyLoD top_lod;
  top_lod.push_back(top_offset);
  top->set_lod(top_lod);
  top->Resize({top_l, _num_emb});
  auto* top_data = dev_ctx.template Alloc<T>(top);

  LegacyLoD drop_pos_lod;
  drop_pos_lod.push_back(drop_pos_offset);
  drop_pos->set_lod(drop_pos_lod);

  iter = dev_ctx.template Alloc<int>(drop_pos);
  int top_counter = 0;
  for (size_t i = 0; i < offset.size() - 1; ++i) {
    int w_drop = static_cast<int>(drop_pos_offset[i + 1] - drop_pos_offset[i]);
    int w = static_cast<int>(offset[i + 1] - offset[i]);
    if (w_drop == 0) {
      if (w >= 2) {
        for (int ilayer = 1; ilayer < _pyramid_layer && ilayer < w; ++ilayer) {
          for (int l = 0; l < w - ilayer; ++l) {
            iter++;
          }
        }
      }
      auto* top_pos = top_data + top_counter++ * _num_emb;
      memset(top_pos, 0, _num_emb * sizeof(T));
      continue;
    }
    if (w >= 2) {
      for (int ilayer = 1; ilayer < _pyramid_layer && ilayer < w; ++ilayer) {
        for (int l = 0; l < w - ilayer; ++l) {
          if (*(iter++) == 0) {
            // do nothing
          } else {
            auto* top_pos = top_data + top_counter++ * _num_emb;
            hash_embedding_ff<T>((const float*)(bottom_data + offset[i] + l),
                                 ilayer + 1,
                                 top_pos,
                                 weights,
                                 _num_emb,
                                 _rand_len,
                                 _space_len);
          }
        }
      }
    }
  }
  if (iter != iter_end) {
    exit(1);
  }
  auto weight_type = _blobs_0->dtype();
  if (_is_training == 0 && weight_type != DataType::INT8) {
    funcs::axpy_noadd(
        top_data, top_data, top->dims()[0] * top->dims()[1], _drop_out_percent);
  }
}
#endif

#ifdef _WIN32
template <typename T, typename Context>
void CPUPyramidHashOPKernel(const Context& dev_ctx,
                            const DenseTensor& x,
                            const DenseTensor& w,
                            const DenseTensor& white_list,
                            const DenseTensor& black_list,
                            int num_emb,
                            int space_len,
                            int pyramid_layer,
                            int rand_len,
                            float drop_out_percent,
                            int is_training,
                            bool use_filter,
                            int white_list_len,
                            int black_list_len,
                            int seed,
                            float lr,
                            const std::string& distribute_update_vars,
                            DenseTensor* out,
                            DenseTensor* drop_pos,
                            DenseTensor* x_temp_out) {}
#endif
}  // namespace phi

PD_REGISTER_KERNEL(
    pyramid_hash, CPU, ALL_LAYOUT, phi::CPUPyramidHashOPKernel, float, int8_t) {
  kernel->InputAt(0).SetDataType(phi::DataType::INT32);
  kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
  kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
}