// Copyright (c) 2019 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 <gtest/gtest.h>
#include <memory>
#include <random>
#include "lite/backends/opencl/target_wrapper.h"
#include "lite/core/op_registry.h"
#include "lite/core/tensor.h"
#include "lite/kernels/opencl/test_helper.h"

#define FP16_MAX_DIFF (5e-1)
namespace paddle {
namespace lite {
void box_coder_ref(float* proposals_data,
                   const float* anchors_data,
                   const float* bbox_deltas_data,
                   const float* variances_data,
                   int axis,
                   bool box_normalized,
                   std::string code_type,
                   int row,
                   int col) {
  if (code_type == "decode_center_size") {
    int anchor_len = 4;
    int out_len = 4;
    int var_len = 4;
    int delta_len = 4;
    float normalized = !box_normalized ? 1.f : 0;

    for (int64_t row_id = 0; row_id < row; ++row_id) {
      for (int64_t col_id = 0; col_id < col; ++col_id) {
        size_t delta_offset = row_id * col * delta_len + col_id * delta_len;
        size_t out_offset = row_id * col * out_len + col_id * out_len;
        int prior_box_offset =
            axis == 0 ? col_id * anchor_len : row_id * anchor_len;
        int var_offset = axis == 0 ? col_id * var_len : row_id * var_len;
        auto anchor_data_tmp = anchors_data + prior_box_offset;
        auto bbox_deltas_data_tmp = bbox_deltas_data + delta_offset;
        auto proposals_data_tmp = proposals_data + out_offset;
        auto anchor_width =
            anchor_data_tmp[2] - anchor_data_tmp[0] + normalized;
        auto anchor_height =
            anchor_data_tmp[3] - anchor_data_tmp[1] + normalized;
        auto anchor_center_x = anchor_data_tmp[0] + 0.5 * anchor_width;
        auto anchor_center_y = anchor_data_tmp[1] + 0.5 * anchor_height;
        float bbox_center_x = 0, bbox_center_y = 0;
        float bbox_width = 0, bbox_height = 0;

        auto variances_data_tmp = variances_data + var_offset;
        bbox_center_x =
            variances_data_tmp[0] * bbox_deltas_data_tmp[0] * anchor_width +
            anchor_center_x;
        bbox_center_y =
            variances_data_tmp[1] * bbox_deltas_data_tmp[1] * anchor_height +
            anchor_center_y;
        bbox_width = std::exp(variances_data_tmp[2] * bbox_deltas_data_tmp[2]) *
                     anchor_width;
        bbox_height =
            std::exp(variances_data_tmp[3] * bbox_deltas_data_tmp[3]) *
            anchor_height;
        proposals_data_tmp[0] = bbox_center_x - bbox_width / 2;
        proposals_data_tmp[1] = bbox_center_y - bbox_height / 2;
        proposals_data_tmp[2] = bbox_center_x + bbox_width / 2 - normalized;
        proposals_data_tmp[3] = bbox_center_y + bbox_height / 2 - normalized;
      }
    }
  } else if (code_type == "encode_center_size") {
    LOG(FATAL) << "not implemented type: " << code_type;
  } else {
    LOG(FATAL) << "not supported type: " << code_type;
  }
}
// #define BOXCODER_FP16_LOOP_TEST
// #define BOXCODER_FP16_PRINT_RESULT
TEST(box_coder_image2d, compute) {
#ifdef BOXCODER_FP16_LOOP_TEST
  for (auto n : {1, 2, 3, 4}) {
    for (auto m : {1, 3, 4, 8}) {
      for (auto norm : {true}) {
        for (auto code_type : {"decode_center_size"}) {
          for (auto axis : {0}) {
#else
  const int n = 1;
  const int m = 1;
  const bool norm = true;
  const std::string code_type = "decode_center_size";
  const int axis = 0;
#endif  // BOXCODER_FP16_LOOP_TEST

            LOG(INFO) << "======== input shape[n,c,h,w]:" << n << " " << m
                      << " ========";
            LOG(INFO) << "======== parameters: norm = " << norm
                      << ", axis = " << axis << "code_type: " << code_type;

            auto kernels = KernelRegistry::Global().Create(
                "box_coder",
                TARGET(kOpenCL),
                PRECISION(kFP16),
                DATALAYOUT(kImageDefault));
            ASSERT_FALSE(kernels.empty());
            auto kernel = std::move(kernels.front());
            LOG(INFO) << "get kernel:" << kernel->doc();

            lite::Tensor prior_box, prior_box_var, target_box, output_box;
            operators::BoxCoderParam param;
            param.prior_box = &prior_box;
            param.prior_box_var = &prior_box_var;
            param.target_box = &target_box;
            param.proposals = &output_box;
            param.axis = axis;
            param.box_normalized = norm;
            param.code_type = code_type;

            std::unique_ptr<KernelContext> context(new KernelContext);
            context->As<OpenCLContext>().InitOnce();

            kernel->SetParam(param);
            std::unique_ptr<KernelContext> boxcoder_context(new KernelContext);
            context->As<OpenCLContext>().CopySharedTo(
                &(boxcoder_context->As<OpenCLContext>()));
            kernel->SetContext(std::move(boxcoder_context));

            const DDim prior_box_dims =
                DDim(std::vector<DDim::value_type>{1, 1, m, 4});
            const DDim prior_box_var_dims =
                DDim(std::vector<DDim::value_type>{1, 1, m, 4});
            const DDim target_box_dims =
                DDim(std::vector<DDim::value_type>{1, n, m, 4});
                  
            const DDim out_dim =
                DDim(std::vector<DDim::value_type>{1, n, m, 4});
            prior_box.Resize(prior_box_dims);
            prior_box_var.Resize(prior_box_var_dims);
            target_box.Resize(target_box_dims);
            output_box.Resize(out_dim);

            std::vector<float> prior_box_data(prior_box_dims.production());
            std::vector<float> prior_box_var_data(prior_box_var_dims.production());
            std::vector<float> target_box_data(target_box_dims.production());
            for (int i = 0; i < prior_box_dims.production(); i++) {
              prior_box_data[i] = i * 1.1 / prior_box_dims.production();
            }
            for (int i = 0; i < prior_box_var_dims.production(); i++) {
              prior_box_var_data[i] = i * 1.2 / prior_box_var_dims.production();
            }
            for (int i = 0; i < target_box_dims.production(); i++) {
              target_box_data[i] = i * 1.3 / target_box_dims.production();
            }

            LOG(INFO) << "prepare input";
            CLImageConverterDefault* default_converter =
                new CLImageConverterDefault();
            DDim prior_box_image_shape =
                default_converter->InitImageDimInfoWith(prior_box_dims);
            LOG(INFO) << "prior_box_image_shape = " << prior_box_image_shape[0] << " "
                      << prior_box_image_shape[1];
            std::vector<half_t> prior_box_image_data(
                prior_box_image_shape.production() * 4);  // 4 : RGBA
            default_converter->NCHWToImage(
                prior_box_data.data(),
                prior_box_image_data.data(),
                prior_box_dims);
            auto* prior_box_image = prior_box.mutable_data<half_t, cl::Image2D>(
                prior_box_image_shape[0],
                prior_box_image_shape[1],
                prior_box_image_data.data());

            DDim prior_box_var_image_shape =
                default_converter->InitImageDimInfoWith(prior_box_var_dims);
            LOG(INFO) << "prior_box_var_image_shape = " << prior_box_var_image_shape[0] << " "
                      << prior_box_var_image_shape[1];
            std::vector<half_t> prior_box_var_image_data(
                prior_box_var_image_shape.production() * 4);  // 4 : RGBA
            default_converter->NCHWToImage(
                prior_box_var_data.data(),
                prior_box_var_image_data.data(),
                prior_box_var_dims);
            auto* prior_box_var_image = prior_box_var.mutable_data<half_t, cl::Image2D>(
                prior_box_var_image_shape[0],
                prior_box_var_image_shape[1],
                prior_box_var_image_data.data());

            DDim target_box_image_shape =
                default_converter->InitImageDimInfoWith(target_box_dims);
            LOG(INFO) << "target_box_image_shape = " << target_box_image_shape[0] << " "
                      << target_box_image_shape[1];
            std::vector<half_t> target_box_image_data(
                target_box_image_shape.production() * 4);  // 4 : RGBA
            default_converter->NCHWToImage(
                target_box_data.data(),
                target_box_image_data.data(),
                target_box_dims);
            auto* target_box_image = target_box.mutable_data<half_t, cl::Image2D>(
                target_box_image_shape[0],
                target_box_image_shape[1],
                target_box_image_data.data());

            DDim out_image_shape =
                default_converter->InitImageDimInfoWith(out_dim);
            LOG(INFO) << "out_image_shape = " << out_image_shape[0] << " "
                      << out_image_shape[1];
            auto* out_image = output_box.mutable_data<half_t, cl::Image2D>(
                out_image_shape[0], out_image_shape[1]);
            kernel->Launch();

            auto* wait_list = context->As<OpenCLContext>().cl_wait_list();
            auto* out_ptr = param.proposals->data<half_t, cl::Image2D>();
            auto it = wait_list->find(out_ptr);
            if (it != wait_list->end()) {
              VLOG(4) << "--- Find the sync event for the target cl "
                        "tensor. ---";
              auto& event = *(it->second);
              event.wait();
            } else {
              LOG(FATAL) << "Could not find the sync event for the "
                                  "target cl tensor.";
            }

            lite::Tensor out_ref_tensor;
            out_ref_tensor.Resize(out_dim);
            box_coder_ref(out_ref_tensor.mutable_data<float>(),
                          prior_box_data.data(),
                          target_box_data.data(),
                          prior_box_var_data.data(),
                          axis,
                          norm,
                          code_type,
                          target_box_dims[0],
                          target_box_dims[1]);

            const size_t cl_image2d_row_pitch{0};
            const size_t cl_image2d_slice_pitch{0};
            half_t* out_image_data =
                new half_t[out_image_shape.production() * 4];
            TargetWrapperCL::ImgcpySync(out_image_data,
                                        out_image,
                                        out_image_shape[0],
                                        out_image_shape[1],
                                        cl_image2d_row_pitch,
                                        cl_image2d_slice_pitch,
                                        IoDirection::DtoH);
            float* out_data = new float[out_image_shape.production() * 4];
            default_converter->ImageToNCHW(
                out_image_data, out_data, out_image_shape, out_dim);
// result
#ifdef BOXCODER_FP16_PRINT_RESULT
            LOG(INFO) << "---- print kernel result (input -> output) ----";
            for (int eidx = 0; eidx < out_dim.production(); ++eidx) {
              std::cout << target_box_data[eidx] << " -> " << out_data[eidx]
                        << std::endl;
            }
#endif  // BOXCODER_FP16_PRINT_RESULT
            const float* out_ref = out_ref_tensor.data<float>();
            for (int i = 0; i < out_dim.production(); i++) {
              auto abs_diff = abs(out_data[i] - out_ref[i]);
              auto relative_diff =
                  COMPUTE_RELATIVE_DIFF(out_data[i], out_ref[i]);
              EXPECT_EQ((relative_diff <= FP16_MAX_DIFF) ||
                            (abs_diff <= FP16_MAX_DIFF),
                            true);
              if ((relative_diff > FP16_MAX_DIFF) &&
                  (abs_diff > FP16_MAX_DIFF)) {
                LOG(ERROR) << "error idx:" << i << ", in_data[" << i
                           << "]: " << target_box_data[i] << ", out_data[" << i
                           << "]: " << out_data[i] << ", out_ref[" << i
                           << "]: " << out_ref[i]
                           << ", abs_diff: " << abs_diff
                           << ", relative_diff: " << relative_diff
                           << ", FP16_MAX_DIFF: " << FP16_MAX_DIFF;
              }
           }
#ifdef BOXCODER_FP16_LOOP_TEST
        }        // axis
      }          // code_type
    }            // norm
  }              // m
}                // n
#else
// nothing to do.
#endif
}

}  // namespace lite
}  // namespace paddle

USE_LITE_KERNEL(box_coder, kOpenCL, kFP16, kImageDefault, ImageDefault);