// 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 "paddle/common/flags.h"
#include "paddle/phi/common/int_array.h"
#include "paddle/phi/core/tensor_utils.h"
#include "paddle/phi/kernels/compare_kernel.h"
#include "paddle/phi/kernels/concat_kernel.h"
#include "paddle/phi/kernels/cumprod_kernel.h"
#include "paddle/phi/kernels/elementwise_divide_kernel.h"
#include "paddle/phi/kernels/elementwise_multiply_kernel.h"
#include "paddle/phi/kernels/flip_kernel.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/funcs/reduce_functor.h"
#include "paddle/phi/kernels/impl/reduce_grad.h"
#include "paddle/phi/kernels/nonzero_kernel.h"
#include "paddle/phi/kernels/prod_grad_kernel.h"
#include "paddle/phi/kernels/reshape_kernel.h"
#include "paddle/phi/kernels/slice_kernel.h"

COMMON_DECLARE_bool(use_accuracy_compatible_kernel);

namespace phi {

// Uses exclusive forward and reverse cumulative products to avoid division
// by zero. See:
// input:                        [    a,     b,     c]
// cumprod(exclusive, normal):   [1    ,     a, a * b]
// cumprod(exclusive, reverse):  [b * c,     c,     1]
// product:                      [b * c, a * c, a * b]
template <typename T, typename Context>
DenseTensor ProdSafeZerosBackward(const Context& dev_ctx,
                                  const DenseTensor& grad,
                                  const DenseTensor& inp,
                                  int dim) {
  if (inp.numel() == 0) {
    DenseTensor result;
    result.Resize(inp.dims());
    dev_ctx.template Alloc<T>(&result);
    return result;
  }

  int64_t dim_size = inp.dims()[dim];
  if (dim_size == 1) {
    DenseTensor result;
    result.Resize(grad.dims());
    dev_ctx.template Alloc<T>(&result);
    Copy(dev_ctx, grad, dev_ctx.GetPlace(), false, &result);
    return result;
  }

  // ones: shape same as inp but with size 1 along dim
  auto ones_dims = vectorize(inp.dims());
  ones_dims[dim] = 1;
  DenseTensor ones =
      Full<T, Context>(dev_ctx, IntArray(ones_dims), static_cast<T>(1));

  // exclusive_normal = cat([ones, inp[:dim_size-1]], dim).cumprod(dim)
  DenseTensor inp_head = Slice<T, Context>(
      dev_ctx, inp, {static_cast<int64_t>(dim)}, {0}, {dim_size - 1});
  DenseTensor exclusive_normal_input =
      Concat<T, Context>(dev_ctx, {&ones, &inp_head}, dim);

  DenseTensor exclusive_normal;
  exclusive_normal.Resize(exclusive_normal_input.dims());
  dev_ctx.template Alloc<T>(&exclusive_normal);
  CumprodKernel<T, Context>(
      dev_ctx, exclusive_normal_input, dim, false, false, &exclusive_normal);

  // exclusive_reverse = cat([ones, flip(inp[1:], dim)], dim).cumprod(dim)
  //                     .flip(dim)
  DenseTensor inp_tail = Slice<T, Context>(
      dev_ctx, inp, {static_cast<int64_t>(dim)}, {1}, {dim_size});

  DenseTensor inp_tail_flipped;
  inp_tail_flipped.Resize(inp_tail.dims());
  dev_ctx.template Alloc<T>(&inp_tail_flipped);
  FlipKernel<T, Context>(dev_ctx, inp_tail, {dim}, &inp_tail_flipped);

  DenseTensor exclusive_reverse_input =
      Concat<T, Context>(dev_ctx, {&ones, &inp_tail_flipped}, dim);

  DenseTensor exclusive_reverse_cumprod;
  exclusive_reverse_cumprod.Resize(exclusive_reverse_input.dims());
  dev_ctx.template Alloc<T>(&exclusive_reverse_cumprod);
  CumprodKernel<T, Context>(dev_ctx,
                            exclusive_reverse_input,
                            dim,
                            false,
                            false,
                            &exclusive_reverse_cumprod);

  DenseTensor exclusive_reverse;
  exclusive_reverse.Resize(exclusive_reverse_cumprod.dims());
  dev_ctx.template Alloc<T>(&exclusive_reverse);
  FlipKernel<T, Context>(
      dev_ctx, exclusive_reverse_cumprod, {dim}, &exclusive_reverse);

  // result = grad * (exclusive_normal * exclusive_reverse)
  DenseTensor product =
      Multiply<T, Context>(dev_ctx, exclusive_normal, exclusive_reverse);
  return Multiply<T, Context>(dev_ctx, grad, product);
}

template <typename T, typename Context>
void ProdGradKernel(const Context& dev_ctx,
                    const DenseTensor& x,
                    const DenseTensor& out,
                    const DenseTensor& out_grad,
                    const IntArray& dims,
                    bool keep_dim,
                    bool reduce_all,
                    DenseTensor* x_grad) {
  if (x_grad && x_grad->numel() == 0) {
    dev_ctx.template Alloc<T>(x_grad);
    return;
  }
  reduce_all = recompute_reduce_all(x, dims, reduce_all);

  if (FLAGS_use_accuracy_compatible_kernel) {
    if (reduce_all) {
      if (x.dims().size() == 0) {
        dev_ctx.template Alloc<T>(x_grad);
        Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
        return;
      }

      // Detect zeros: create (x == 0) mask and count via nonzero
      DenseTensor zeros_like_x = Full<T, Context>(
          dev_ctx, IntArray(vectorize(x.dims())), static_cast<T>(0));

      DenseTensor eq_mask;
      eq_mask.Resize(x.dims());
      dev_ctx.template Alloc<bool>(&eq_mask);
      EqualKernel<T, Context>(dev_ctx, x, zeros_like_x, &eq_mask);

      DenseTensor zero_indices;
      NonZeroKernel<bool, Context>(dev_ctx, eq_mask, &zero_indices);
      int64_t num_zeros = zero_indices.dims()[0];

      dev_ctx.template Alloc<T>(x_grad);

      if (num_zeros == 0) {
        // No zeros: grad * (result / input)
        DenseTensor result_div_input = Divide<T, Context>(dev_ctx, out, x);
        DenseTensor grad_result =
            Multiply<T, Context>(dev_ctx, out_grad, result_div_input);
        Copy(dev_ctx, grad_result, dev_ctx.GetPlace(), false, x_grad);
      } else if (num_zeros > 1) {
        // More than one zero: gradient is all zeros
        FullLikeKernel<T, Context>(dev_ctx,
                                   x,
                                   Scalar(static_cast<T>(0)),
                                   CppTypeToDataType<T>::Type(),
                                   x_grad);
      } else {
        // Exactly one zero (or meta tensor): use safe cumprod backward
        // Flatten to 1D, apply along dim=0, reshape back
        DenseTensor x_flat =
            Reshape<T, Context>(dev_ctx, x, {static_cast<int64_t>(x.numel())});
        DenseTensor grad_result =
            ProdSafeZerosBackward<T, Context>(dev_ctx, out_grad, x_flat, 0);
        auto x_shape = vectorize<int64_t>(x.dims());
        Reshape<T, Context>(dev_ctx, grad_result, x_shape, x_grad);
      }
    } else {
      auto dim_vec = dims.GetData();

      if (x.dims().size() == 0) {
        dev_ctx.template Alloc<T>(x_grad);
        Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
        return;
      }

      if (dim_vec.size() == 1) {
        int64_t dim = dim_vec[0];
        if (dim < 0) {
          dim += x.dims().size();
        }

        // Unsqueeze grad and result if !keepdim, to match input rank
        DenseTensor grad_expanded = out_grad;
        DenseTensor result_expanded = out;
        if (!keep_dim) {
          auto grad_shape = vectorize<int64_t>(out_grad.dims());
          grad_shape.insert(grad_shape.begin() + dim, 1);
          grad_expanded = Reshape<T, Context>(dev_ctx, out_grad, grad_shape);

          auto result_shape = vectorize<int64_t>(out.dims());
          result_shape.insert(result_shape.begin() + dim, 1);
          result_expanded = Reshape<T, Context>(dev_ctx, out, result_shape);
        }

        // Detect zeros
        DenseTensor zeros_like_x = Full<T, Context>(
            dev_ctx, IntArray(vectorize(x.dims())), static_cast<T>(0));

        DenseTensor eq_mask;
        eq_mask.Resize(x.dims());
        dev_ctx.template Alloc<bool>(&eq_mask);
        EqualKernel<T, Context>(dev_ctx, x, zeros_like_x, &eq_mask);

        DenseTensor zero_indices;
        NonZeroKernel<bool, Context>(dev_ctx, eq_mask, &zero_indices);
        int64_t total_zeros = zero_indices.dims()[0];

        dev_ctx.template Alloc<T>(x_grad);

        if (total_zeros == 0) {
          // No zeros: grad * (result / input) with broadcasting
          DenseTensor result_div_input =
              Divide<T, Context>(dev_ctx, result_expanded, x);
          DenseTensor grad_result =
              Multiply<T, Context>(dev_ctx, grad_expanded, result_div_input);
          Copy(dev_ctx, grad_result, dev_ctx.GetPlace(), false, x_grad);
        } else {
          // Has zeros: use safe cumprod backward
          DenseTensor grad_result = ProdSafeZerosBackward<T, Context>(
              dev_ctx, grad_expanded, x, static_cast<int>(dim));
          Copy(dev_ctx, grad_result, dev_ctx.GetPlace(), false, x_grad);
        }
      } else {
        // Multiple dims: fall back to original Eigen-based implementation
        ReduceGradKernel<Context, T, funcs::ProdGradFunctor>(dev_ctx,
                                                             x,
                                                             out,
                                                             out_grad,
                                                             dims.GetData(),
                                                             keep_dim,
                                                             reduce_all,
                                                             x_grad);
      }
    }
  } else {
    ReduceGradKernel<Context, T, funcs::ProdGradFunctor>(dev_ctx,
                                                         x,
                                                         out,
                                                         out_grad,
                                                         dims.GetData(),
                                                         keep_dim,
                                                         reduce_all,
                                                         x_grad);
  }
}

}  // namespace phi