// Copyright (c) 2023 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 <glog/logging.h>
#include <gtest/gtest.h>
#include <memory>
#include <sstream>
#include <string>
#include <tuple>
#include <unordered_map>

#include "paddle/cinn/hlir/dialect/operator/ir/cinn_op.h"
#include "paddle/cinn/hlir/dialect/operator/ir/op_attribute.h"
#include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
#include "paddle/cinn/hlir/dialect/runtime/ir/jit_kernel_op.h"
#include "paddle/cinn/hlir/dialect/runtime/ir/runtime_dialect.h"
#include "paddle/cinn/hlir/framework/pir/utils.h"
#include "paddle/cinn/hlir/framework/pir_compiler.h"
#include "paddle/cinn/utils/data_util.h"
#include "paddle/fluid/framework/new_executor/interpretercore.h"
#include "paddle/fluid/pir/dialect/kernel/ir/kernel_dialect.h"
#include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
#include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
#include "paddle/fluid/pir/transforms/pd_op_to_kernel_pass.h"
#include "paddle/pir/include/core/ir_context.h"
#include "paddle/pir/include/core/program.h"
#include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"

using cinn::hlir::framework::pir::CompatibleInfo;
using cinn::hlir::framework::pir::OpLoweringGroup;
using cinn::hlir::framework::pir::OpLoweringGroupPtr;

bool simple_cmp(float a, float b) { return std::abs((a - b) / a) < 1e-5; }
using ProgramInfo = std::tuple<std::shared_ptr<::pir::Program>,
                               std::vector<OpLoweringGroupPtr>>;
ProgramInfo BuildProgram() {
  ::pir::IrContext* ctx = ::pir::IrContext::Instance();
  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
  auto program = std::make_shared<::pir::Program>(ctx);
  ::pir::Builder builder = ::pir::Builder(ctx, program->block());

  const float value_one = 1.0;  // relu(tan(1.)) = 1.5;
  const float value_two = 2.0;  // relu(tan(2.)) = 0.
  auto full_op_x =
      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64, 128},
                                             value_one,
                                             phi::DataType::FLOAT32,
                                             phi::GPUPlace());

  auto full_op_y =
      builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64, 128},
                                             value_two,
                                             phi::DataType::FLOAT32,
                                             phi::GPUPlace());

  auto tan_op_x = builder.Build<paddle::dialect::TanOp>(full_op_x->result(0));
  auto relu_op_x = builder.Build<paddle::dialect::ReluOp>(tan_op_x->result(0));
  auto tan_op_y = builder.Build<paddle::dialect::TanOp>(relu_op_x->result(0));
  auto relu_op_y = builder.Build<paddle::dialect::ReluOp>(tan_op_y->result(0));

  builder.Build<pir::YieldOp>(std::vector<pir::Value>{full_op_x.result(0)});
  builder.Build<pir::YieldOp>(std::vector<pir::Value>{full_op_y.result(0)});
  builder.Build<pir::YieldOp>(std::vector<pir::Value>{relu_op_y.result(0)});

  std::vector<OpLoweringGroupPtr> groups;
  const auto full_op_x_ops =
      std::initializer_list<::pir::Operation*>({full_op_x.operation()});
  groups.emplace_back(std::make_shared<OpLoweringGroup>(
      full_op_x_ops,
      CompatibleInfo::GroupOpsName(full_op_x_ops)));  // For coverage
  groups[0]->mut_output_values().push_back(groups[0]->ops().back()->result(0));

  const auto full_op_y_ops =
      std::initializer_list<::pir::Operation*>({full_op_x.operation()});
  groups.emplace_back(std::make_shared<OpLoweringGroup>(
      full_op_y_ops, CompatibleInfo::GroupOpsName(full_op_y_ops)));

  groups[1]->mut_output_values().push_back(groups[1]->ops().back()->result(0));
  const auto vector_ops =
      std::vector<::pir::Operation*>({tan_op_x.operation(),
                                      relu_op_x.operation(),
                                      tan_op_y.operation(),
                                      relu_op_y.operation()});
  groups.emplace_back(std::make_shared<OpLoweringGroup>(
      vector_ops, CompatibleInfo::GroupOpsName(vector_ops)));
  groups[2]->mut_output_values().push_back(groups[2]->ops().back()->result(0));

  return {program, groups};
}

ProgramInfo BuildSoftmax() {
  ::pir::IrContext* ctx = ::pir::IrContext::Instance();
  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
  ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
  auto program = std::make_shared<::pir::Program>(ctx);
  ::pir::Builder builder = ::pir::Builder(ctx, program->block());
  std::vector<int64_t> axes{-1};

  auto x = builder
               .Build<paddle::dialect::FullOp>(std::vector<int64_t>({16, 16}),
                                               1.0,
                                               phi::DataType::FLOAT32,
                                               phi::GPUPlace(0))
               .result(0);
  auto max = builder.Build<cinn::dialect::ReduceMaxOp>(x, axes, true).result(0);
  auto broadcast_1 =
      builder
          .Build<cinn::dialect::BroadcastOp>(
              max, std::vector<int64_t>({0, 1}), std::vector<int64_t>({16, 16}))
          .result(0);
  auto sub =
      builder.Build<paddle::dialect::SubtractOp>(x, broadcast_1).result(0);
  auto exp = builder.Build<paddle::dialect::ExpOp>(sub).result(0);
  auto sum =
      builder.Build<cinn::dialect::ReduceSumOp>(exp, axes, true).result(0);

  auto broadcast_2 =
      builder
          .Build<cinn::dialect::BroadcastOp>(
              sum, std::vector<int64_t>({0, 1}), std::vector<int64_t>({16, 16}))
          .result(0);
  auto divide =
      builder.Build<paddle::dialect::DivideOp>(exp, broadcast_2).result(0);
  auto yield_op = builder.Build<pir::YieldOp>(std::vector<pir::Value>{divide});

  std::vector<OpLoweringGroupPtr> groups;
  const auto vector_ops =
      std::initializer_list<::pir::Operation*>({max.defining_op(),
                                                broadcast_1.defining_op(),
                                                sub.defining_op(),
                                                exp.defining_op(),
                                                sum.defining_op(),
                                                broadcast_2.defining_op(),
                                                divide.defining_op()});
  groups.emplace_back(std::make_shared<OpLoweringGroup>(
      vector_ops, CompatibleInfo::GroupOpsName(vector_ops)));
  groups[0]->mut_output_values().push_back(groups[0]->ops().back()->result(0));
  groups[0]->set_op_pattern_kind(cinn::hlir::framework::kReduction);

  return {program, groups};
}

// TEST(PirCompiler, CompileSoftmax) {
//   // Step 1: Construct pir::Program
//   ::pir::IrContext* ctx = ::pir::IrContext::Instance();
//   ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
//   ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
//   ctx->GetOrRegisterDialect<cinn::dialect::RuntimeDialect>();
//   ctx->GetOrRegisterDialect<paddle::dialect::KernelDialect>();
//   auto new_program = std::make_shared<::pir::Program>(ctx);

//   auto prog_info = BuildSoftmax();
//   std::shared_ptr<::pir::Program> program = std::get<0>(prog_info);
//   std::vector<GroupPtr> groups = std::get<1>(prog_info);
//   EXPECT_EQ(program->block()->size(), 9u);
//   LOG(INFO) << program->block()->size();

//   std::stringstream ss;
//   program->Print(ss);
//   LOG(INFO) << ss.str();

//   // Step 2: Compiler New pir::Program into Runtime Program
//   auto target = cinn::common::DefaultNVGPUTarget();
//   auto scope = cinn::hlir::framework::BuildScope(target, *program);
//   LOG(INFO) << scope->var_names().size();
//   ASSERT_EQ(scope->var_names().size(), 8);

//   cinn::hlir::framework::PirCompiler ir_compiler(*program, target, scope);
//   auto fn_ptr_res = ir_compiler.BuildCUDAJITInfo(groups);

//   ::pir::Builder builder = ::pir::Builder(ctx, new_program->block());
//   auto x = builder
//                .Build<paddle::dialect::FullOp>(std::vector<int64_t>({16,
//                16}),
//                                                1.0,
//                                                phi::DataType::FLOAT32,
//                                                phi::GPUPlace(0))
//                .result(0);

//   std::unordered_map<std::string, ::pir::Attribute> op_attrs{
//       {cinn::dialect::JitKernelOp::kAttrName,
//        cinn::dialect::CINNKernelInfoAttribute::get(ctx, fn_ptr_res[0])},
//   };

//   std::vector<pir::Type> vec_types;

//   vec_types.push_back(groups[0]->ops.back()->result(0).type());

//   std::string jit_op_name = cinn::dialect::JitKernelOp::name();
//   ::pir::OpInfo op_info = ctx->GetRegisteredOpInfo(jit_op_name);
//   ::pir::Operation* cinn_op =
//       ::pir::Operation::Create({x}, op_attrs, vec_types, op_info);

//   new_program->block()->push_back(cinn_op);

//   builder.SetInsertionPointToBlockEnd(new_program->block());
//   builder.Build<paddle::dialect::FetchOp>(
//       cinn_op->result(cinn_op->num_results() - 1), "out", 0);

//   phi::Place place = phi::GPUPlace(0);

//   auto kernel_program =
//       pir::PdOpLowerToKernelPass(new_program.get(), place);

//   paddle::framework::Scope exe_scope;

//   paddle::framework::interpreter::ExecutionConfig exe_conf;
//   exe_conf.create_local_scope = false;
//   paddle::framework::InterpreterCore executor(
//       place, {"out@fetch"}, kernel_program->block(), &exe_scope);

//   executor.Run({}, true);
//   auto out_tensor =
//       executor.local_scope()->FindVar("out@fetch")->Get<phi::DenseTensor>();
//   bool res0 = simple_cmp(out_tensor.data<float>()[0], 1.0 / 16);
//   EXPECT_EQ(res0, true);
// }

// TEST(PirCompiler, CompileGroupOps) {
//   // Step 1: Construct pir::Program
//   auto prog_info = BuildProgram();
//   std::shared_ptr<::pir::Program> program = std::get<0>(prog_info);
//   std::vector<GroupPtr> groups = std::get<1>(prog_info);
//   EXPECT_EQ(program->block()->size(), 9u);
//   LOG(INFO) << program->block()->size();

//   std::stringstream ss;
//   program->Print(ss);
//   LOG(INFO) << ss.str();

//   // Step 2: Compiler New pir::Program into Runtime Program
//   auto target = cinn::common::DefaultNVGPUTarget();
//   auto scope = cinn::hlir::framework::BuildScope(target, *program);
//   ASSERT_EQ(scope->var_names().size(), 6);

//   cinn::hlir::framework::PirCompiler ir_compiler(*program, target, scope);
//   auto runtime_program = ir_compiler.Build(groups);

//   // Step 3: Execute Runtime Instruction and check Scope.
//   ASSERT_NO_THROW(runtime_program->Execute());
//   for (auto& var_name : scope->var_names()) {
//     std::string name = {var_name.begin(), var_name.end()};
//     std::vector<float> data =
//         cinn::GetTensorData<float>(scope->GetTensor(name), target);
//     for (int i = 0; i < 1; ++i) {
//       LOG_FIRST_N(INFO, 10) << "data: " << data[i];
//     }
//   }
// }