// Copyright (c) 2022 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 "glog/logging.h"

#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/kernels/complex_kernel.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/impl/einsum_kernel_impl.h"
#include "paddle/phi/kernels/tile_grad_kernel.h"
#include "paddle/phi/kernels/tile_kernel.h"
#include "paddle/utils/string/string_helper.h"
namespace phi {

template <typename T, typename Context>
DenseTensor PerformTileAndReduction(const Context& dev_ctx,
                                    const LabelMap& label2type,
                                    const LabelMap& label2shape,
                                    const std::vector<int64_t>& broadcast_shape,
                                    const std::vector<int64_t> x_shape,
                                    std::string equ,   // value pass
                                    DenseTensor& t) {  // NOLINT
  auto tmp_label = equ;
  auto tmp_union = unique_labels(tmp_label);
  auto op_label = std::string(tmp_union.begin(), tmp_union.end());
  VLOG(5) << "Start PerformTileAndReduction equation " << equ
          << " with operand shape: "
          << paddle::string::join_strings(vectorize<int64_t>(t.dims()), ",");
  DenseTensor ret;
  std::vector<int64_t> repeat_times;
  std::vector<int64_t> resize_dims;
  std::vector<int64_t> recover_shape;
  std::vector<int64_t> t_shape = vectorize<int64_t>(t.dims());
  for (size_t i = 0; i < op_label.size(); i++) {
    int c = op_label[i];
    if (label2type[c] == LabelType::Reduction) {
      repeat_times.push_back(label2shape[c]);
      resize_dims.push_back(1);
      recover_shape.push_back(label2shape[c]);
      t_shape.insert(t_shape.begin() + i, 1);
    } else {
      resize_dims.push_back(label2shape[c]);
      repeat_times.push_back(1);
      recover_shape.push_back(label2shape[c]);
    }
  }
  PADDLE_ENFORCE_EQ(op_label.size(),
                    t_shape.size(),
                    common::errors::InvalidArgument(
                        "Input shape size doesn't match label nums, input "
                        "shape size: `%d`, but got label nums: `%d`",
                        t_shape.size(),
                        op_label.size()));
  for (size_t i = 0; i < op_label.size(); i++) {
    int c = op_label[i];
    if (label2type[c] == LabelType::Contraction &&
        t_shape[i] != label2shape[c]) {
      repeat_times[i] = label2shape[c];
      resize_dims[i] = 1;
    }
  }
  t.Resize(resize_dims);
  DenseTensor after_tile;
  if (std::all_of(repeat_times.begin(), repeat_times.end(), [](int64_t x) {
        return x == 1;
      })) {
    after_tile = t;
  } else {
    VLOG(4) << "do TileKernel with repeat_times="
            << paddle::string::join_strings(repeat_times, ",");
    TileKernel<T, Context>(dev_ctx, t, repeat_times, &after_tile);
  }
  ret = after_tile;
  VLOG(5) << "PermformTileAndReduction: recover shape: "
          << paddle::string::join_strings(recover_shape, ",");
  ret.Resize(recover_shape);

  // undiagonalize by einsum equation. only contain undiagonal operations.
  DenseTensor undiagonal_out;
  if (op_label != equ) {
    VLOG(5) << "Undiagonal by einsum with args: " << op_label + "->" + equ;
    EinsumInferKernel<T, Context>(
        dev_ctx, {&ret}, op_label + "->" + equ, &undiagonal_out);
  } else {
    undiagonal_out = ret;
  }

  // call TileGradKernel to reverse broadcast operation.
  VLOG(5) << "After diagonalize, we have tensor with shape: "
          << paddle::string::join_strings(
                 vectorize<int64_t>(undiagonal_out.dims()), ',');
  repeat_times.clear();
  for (size_t i = 0; i < x_shape.size(); ++i) {
    VLOG(4) << "broadcast shape is " << broadcast_shape[i] << ", x_shape is "
            << x_shape[i];
    repeat_times.push_back(broadcast_shape[i] / x_shape[i]);
  }
  bool is_all_ones = std::all_of(repeat_times.begin(),
                                 repeat_times.end(),
                                 [](int64_t x) { return x == 1; });
  if (is_all_ones) {
    VLOG(4) << "don't need broadcast recover, we just return undiagonal_out.";
    return undiagonal_out;
  }
  DenseTensor tmp_x;
  DenseTensor broadcast_out;
  tmp_x.Resize(x_shape);
  broadcast_out.Resize(x_shape);
  TileGradKernel<T, Context>(
      dev_ctx, tmp_x, undiagonal_out, repeat_times, &broadcast_out);
  VLOG(5) << "After broadcast recover, we have tensor with shape: "
          << paddle::string::join_strings(
                 vectorize<int64_t>(broadcast_out.dims()), ',');
  return broadcast_out;
}

template <typename T, typename Context>
void EinsumGradKernel(const Context& dev_ctx,
                      const std::vector<const DenseTensor*>& x,
                      const std::vector<const DenseTensor*>& inner_cache,
                      const DenseTensor& out_grad,
                      const std::string& equation,
                      std::vector<DenseTensor*> x_grad) {
  VLOG(5) << "Start EinsumGradKernel:";
  bool has_zero_size_tensor = out_grad.numel() == 0;
  for (auto& i : x_grad) {
    if (i != nullptr) {
      if (i->numel() == 0) {
        has_zero_size_tensor = true;
      }
      Full<T, Context>(dev_ctx, i->dims(), 0, i);
    }
  }
  if (has_zero_size_tensor) return;
  LabelMap labelshape(0);
  LabelMap labeltype(LabelType::Reduction);
  std::vector<LabelMap> label2perms(x.size(), LabelMap(-1));
  std::vector<char> all_labels;  // order: ABO, AO, BO, AB, Reduce
  std::vector<std::vector<int64_t>> broadcast_shapes(2);
  std::vector<int64_t> output_dims;

  std::vector<DDim> input_dims;
  for (auto& i : x) {
    input_dims.push_back(i->dims());
  }
  std::vector<std::string> input_strs;
  std::string right;
  ParseEinsumEquation(equation,
                      input_dims,
                      &labelshape,
                      &labeltype,
                      &all_labels,
                      &label2perms,
                      &broadcast_shapes,
                      &output_dims,
                      &right,
                      &input_strs);

  VLOG(4) << "After grad parse einsum equation.";

  auto gather_labels_except_reduction = [&labeltype](std::string all) {
    std::string res("");
    for (auto c : all)
      if (labeltype[static_cast<int>(c)] != LabelType::Reduction) res += c;
    auto tmp_unique = unique_labels(res);
    return std::string(tmp_unique.begin(), tmp_unique.end());
  };
  if (x.size() == 1) {  // Unary
    auto splits = paddle::string::split_string(equation, "->");
    auto left = splits[0];
    right = splits[1];
    auto new_equation = right + "->" + gather_labels_except_reduction(left);
    auto new_operands = std::vector<const DenseTensor*>();
    new_operands.push_back(&out_grad);
    DenseTensor before_tile;
    VLOG(5) << "new_equation is " << new_equation;
    EinsumInferKernel<T, Context>(
        dev_ctx, new_operands, new_equation, &before_tile);
    *(x_grad[0]) =
        PerformTileAndReduction<T, Context>(dev_ctx,
                                            labeltype,
                                            labelshape,
                                            broadcast_shapes[0],
                                            vectorize<int64_t>(x[0]->dims()),
                                            left,
                                            before_tile);
#ifndef PADDLE_WITH_XPU  // xpu is not support conj now, we just disable it.
    *(x_grad[0]) = Conj<T, Context>(dev_ctx, *x_grad[0]);
#endif
  } else {
    auto splits = paddle::string::split_string(equation, "->");
    auto left = splits[0];
    auto ops = paddle::string::split_string(left, ",");
    right = splits[1];
    auto equation_for_A =
        ops[1] + "," + right + "->" + gather_labels_except_reduction(ops[0]);
    auto equation_for_B =
        right + "," + ops[0] + "->" + gather_labels_except_reduction(ops[1]);
    auto operands_for_A = std::vector<const DenseTensor*>();
    auto operands_for_B = std::vector<const DenseTensor*>();
    DenseTensor dA, dB;
#ifndef PADDLE_WITH_XPU  // xpu is not support conj now, we just disable it.
    auto out_grad_conj = Conj<T, Context>(dev_ctx, out_grad);
#else
    auto out_grad_conj = out_grad;
#endif
    // dA = einsum(B, dC)
    operands_for_A.push_back(x[1]);
    operands_for_A.push_back(&out_grad_conj);
    // dB = einsum(dC, A)
    operands_for_B.push_back(&out_grad_conj);
    operands_for_B.push_back(x[0]);

    std::vector<DenseTensor> cache(3);  // set empty; TA, TB, TdC
    if (inner_cache.size() >
        0) {  // for compatibility,  we can load and run v2.3 EinsumOp.
      cache[0].ShareBufferWith(*(inner_cache[0]));
      cache[1].ShareBufferWith(*(inner_cache[1]));
    }
    EinsumKernelImpl<T, Context>(dev_ctx,
                                 all_labels,
                                 labelshape,
                                 operands_for_A,
                                 equation_for_A,
                                 &dA,
                                 {&cache[1], &cache[2]},
                                 false);

    EinsumKernelImpl<T, Context>(dev_ctx,
                                 all_labels,
                                 labelshape,
                                 operands_for_B,
                                 equation_for_B,
                                 &dB,
                                 {&cache[2], &cache[0]},
                                 false);

    // release the cache tensor dTC to save memory right now. they are useless
    // now.
    cache.clear();
    if (x_grad[0]) {
      *(x_grad[0]) =
          PerformTileAndReduction<T, Context>(dev_ctx,
                                              labeltype,
                                              labelshape,
                                              broadcast_shapes[0],
                                              vectorize<int64_t>(x[0]->dims()),
                                              ops[0],
                                              dA);
      VLOG(4) << "After call dA";
#ifndef PADDLE_WITH_XPU  // xpu is not support conj now, we just disable it.
      *(x_grad[0]) = Conj<T, Context>(dev_ctx, *x_grad[0]);
#endif
    }
    if (x_grad[1]) {
      *(x_grad[1]) =
          PerformTileAndReduction<T, Context>(dev_ctx,
                                              labeltype,
                                              labelshape,
                                              broadcast_shapes[1],
                                              vectorize<int64_t>(x[1]->dims()),
                                              ops[1],
                                              dB);
#ifndef PADDLE_WITH_XPU  // xpu is not support conj now, we just disable it.
      *(x_grad[1]) = Conj<T, Context>(dev_ctx, *x_grad[1]);
#endif
      VLOG(4) << "After call dA";
    }
  }
}
}  // namespace phi