/* Copyright (c) 2021 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/common/ddim.h>
#include <string>
#include <vector>
#include "paddle/common/flags.h"

namespace phi {

namespace funcs {
inline bool CheckIsLastDimsMatch(const DDim& first, const DDim& second) {
  auto n1 = first.size();
  auto n2 = second.size();
  size_t min_len = std::min(n1, n2);

  for (size_t i = 0; i < min_len; i++) {
    if (first[n1 - 1 - i] != second[n2 - 1 - i]) {
      return false;
    }
  }
  return true;
}
// check whether the tensor with dimension of second can assign to the
// tensor with dimension of first
inline bool CheckIsDimsMatchBool(const DDim& first, const DDim& second) {
  int ignore_axis1 = 0, ignore_axis2 = 0;
  for (; ignore_axis1 < first.size(); ++ignore_axis1) {
    if (first[ignore_axis1] != 1) {
      break;
    }
  }
  for (; ignore_axis2 < second.size(); ++ignore_axis2) {
    if (second[ignore_axis2] != 1) {
      break;
    }
  }

  if (second.size() == ignore_axis2) {
    // second tensor has only one value
    return true;
  }

  if (first.size() - ignore_axis1 >= second.size() - ignore_axis2) {
    auto idx1 = first.size() - 1;
    auto idx2 = second.size() - 1;
    bool is_match = true;
    for (; idx2 >= ignore_axis2; idx2--) {
      if (first[idx1--] != second[idx2] && second[idx2] != 1) {
        is_match = false;
        break;
      }
    }
    if (is_match) {
      return true;
    }
  }

  return false;
}

inline void CheckIsDimsMatch(const DDim& first, const DDim& second) {
  if (CheckIsDimsMatchBool(first, second)) {
    return;
  }
  PADDLE_THROW(errors::InvalidArgument(
      "The shape of tensor assigned value must match the shape "
      "of target shape: %d, but now shape is %d.",
      second.to_str(),
      first.to_str()));
}
/**
 * @brief Normalizes the slice interval [st, ed) with a given step and dimension
 * size.
 *
 * This function adjusts the interval [st, ed) to fit within the bounds defined
 * by the dimension size, taking into account the specified step. It handles
 * both positive and negative steps and accounts for negative indices by
 * converting them to equivalent positive indices within the dimension size.
 *
 * @tparam T The data type of the input parameters, which can be an integer or
 * floating-point type.
 * @param st The starting index of the interval.
 * @param ed The ending index of the interval (exclusive).
 * @param step The step size for iterating through the interval, which can be
 * positive or negative.
 * @param dim_size The size of the dimension, serving as the upper bound for
 * valid indices.
 * @param st_out Pointer to store the normalized starting index.
 * @param ed_out Pointer to store the normalized ending index.
 * @param zero_dim_out Pointer to a boolean flag that is set to true if the
 * resulting interval is empty.
 *
 * @details
 * - If `step > 0`, the function ensures that `st` and `ed` are adjusted to be
 * within the range [0, dim_size).
 * - If `step < 0`, the function adjusts `st` and `ed` to accommodate the
 * reverse traversal of the interval.
 * - Handles special cases where `st` and `ed` may be out of bounds or where
 * `dim_size` is zero.
 * - Uses pointer parameters for output to modify the values directly.
 * - The function also handles scenarios involving negative indices, converting
 * them appropriately.
 *
 * @example
 * T st_out, ed_out;
 * bool zero_dim;
 * normalize_interval(-3, -2, 1, 4, &st_out, &ed_out, &zero_dim);
 * // Results in: st_out = 1, ed_out = 2, zero_dim = false
 *
 * @note The function assumes that the pointers provided for output parameters
 * are valid and non-null.
 */
template <typename T>
void normalize_interval(
    T st, T ed, T step, T dim_size, T* st_out, T* ed_out, bool* zero_dim_out) {
  /* Normalize slice interval [st, ed) with given step and dim_size.
  e.g. if given st = -3, ed = -2, step = 1, dim_size = 4,
  then normalized st_out = 1(-3+4), st_ed = 2(-2+4).

  This function is general enough and applicable
  for both step > 0 and step < 0 scenarios.

  Indicices dipicted as below:

  ===============================================================
                |  0   1     2     3    ...  D-1 | D D+1 ...
  ... -D-2 -D-1 | -D  -D+1  -D+2  -D+3  ... -1   |
  ===============================================================
  */
  // 0 dim size, just return
  if (dim_size <= 0) {
    *st_out = *ed_out = 0;
    *zero_dim_out = true;
    return;
  }

  if (step > 0) {
    /* positive step */
    // 0 dim size case 1
    if (st >= dim_size) {
      *st_out = *ed_out = 0;
      *zero_dim_out = true;
      return;
    }

    // 0 dim size case 2
    if (ed <= -dim_size) {
      *st_out = *ed_out = 0;
      *zero_dim_out = true;
      return;
    }

    // make st belongs: (-inf, -D-1)∪[0, D)
    if (-dim_size <= st && st < 0) {
      st += dim_size;
    }
    // make st belongs: [0, D)
    st = std::max(st, static_cast<T>(0));

    // make ed belongs: [0, +inf)
    if (-dim_size <= ed && ed < 0) {
      ed += dim_size;
    }
    // make ed belongs: [0, D]
    ed = std::min(ed, dim_size);

    // 0 dim size case 3
    if (st >= ed) {
      *st_out = *ed_out = 0;
      *zero_dim_out = true;
      return;
    }
    *st_out = st;
    *ed_out = ed;
    return;

  } else {
    /* negative step */
    // 0 dim size case 1
    if (st <= -dim_size - 1) {
      *st_out = *ed_out = 0;
      *zero_dim_out = true;
      return;
    }

    // 0 dim size case 2
    if (ed >= dim_size - 1) {
      *st_out = *ed_out = 0;
      *zero_dim_out = true;
      return;
    }

    // make st belongs: [0, D)∪[0, +inf)
    if (-dim_size <= st && st < 0) {
      st += dim_size;
    }
    // make st belongs: [0, D)
    st = std::min(st, dim_size - 1);

    // make ed belongs: [-inf, -D)∪[0, D)
    if (-dim_size <= ed && ed < 0) {
      ed += dim_size;
    }
    // make ed belongs: [-D-1, -D)∪[0, D) ==> {-D-1}∪[0, D)
    ed = std::max(ed, -dim_size - 1);

    if (ed == -dim_size - 1) {
      // When ed=-D-1, it is symmetrical to when step is greater than 0 and
      // ed=D.
      *st_out = st;
      *ed_out = ed;
      return;
    }

    // now only remain the case that ed belongs to: [0, D)
    // 0 dim size case 3
    if (ed >= st) {
      *st_out = *ed_out = 0;
      *zero_dim_out = true;
      return;
    }

    *st_out = st;
    *ed_out = ed;
    return;
  }
}

template <typename T = int64_t>
inline void CheckAndUpdateSliceAttrs(const DDim in_dims,
                                     const std::vector<T>& axes,
                                     std::vector<T>* starts,
                                     std::vector<T>* ends,
                                     std::vector<int64_t>* steps = nullptr,
                                     std::vector<T>* infer_flags = nullptr) {
  for (size_t i = 0; i < axes.size(); ++i) {
    T axis = axes[i];
    PADDLE_ENFORCE_LT(
        axis,
        in_dims.size(),
        common::errors::InvalidArgument(
            "The axis value should be less than the rank of input, "
            "but received axes[%d] = %d, rank of input is %d.",
            i,
            axis,
            in_dims.size()));

    if (infer_flags != nullptr && (*infer_flags)[i] == -1) {
      continue;
    }

    T dim_value = in_dims[axis];

    if (dim_value > 0) {
      T step = steps == nullptr ? 1 : (*steps)[i];
      PADDLE_ENFORCE_NE(
          step,
          0,
          common::errors::InvalidArgument(
              "Step should not be 0, but received step = %d.", step));

      T start, end;
      bool dummy_zero_out_dim = false;
      normalize_interval((*starts)[i],
                         (*ends)[i],
                         step,
                         dim_value,
                         &start,
                         &end,
                         &dummy_zero_out_dim);
      if (end == -dim_value - 1) {
        end = -1;
      }

      (*starts)[i] = start;
      (*ends)[i] = end;
    } else if (dim_value == 0) {
      (*starts)[i] = 0;
      (*ends)[i] = 0;
    }
  }
}

template <typename T = int64_t>
inline void UpdateSliceAttrs(const DDim in_dims,
                             const std::vector<T>& axes,
                             std::vector<T>* starts,
                             std::vector<T>* ends,
                             std::vector<int64_t>* steps = nullptr,
                             std::vector<T>* infer_flags = nullptr) {
  for (size_t i = 0; i < axes.size(); ++i) {
    T axis = axes[i];
    if (infer_flags != nullptr && (*infer_flags)[i] == -1) {
      continue;
    }
    T dim_value = in_dims[axis];
    if (dim_value > 0) {
      T step = steps == nullptr ? 1 : (*steps)[i];
      T start = (*starts)[i];
      T end = (*ends)[i];

      bool dummy_zero_out_dim = false;
      normalize_interval(
          start, end, step, dim_value, &start, &end, &dummy_zero_out_dim);

      // manually set the end to -1 when step < 0,
      // which indicates that it can extend to the left endpoint.
      if (end == -dim_value - 1 && step < 0) {
        end = -1;
      }
      (*starts)[i] = start;
      (*ends)[i] = end;
    } else if (dim_value == 0) {
      (*starts)[i] = 0;
      (*ends)[i] = 0;
    }
  }
}

template <typename T = int64_t>
inline DDim GetSliceDims(const DDim in_dims,
                         const std::vector<T>& axes,
                         const std::vector<T>& starts,
                         const std::vector<T>& ends,
                         std::vector<T>* steps = nullptr,
                         std::vector<T>* infer_flags = nullptr) {
  DDim slice_dims(in_dims);

  for (size_t i = 0; i < axes.size(); ++i) {
    T axis = axes[i];
    if (infer_flags != nullptr && (*infer_flags)[i] == -1) {
      slice_dims[axis] = -1;
      continue;
    }

    if (in_dims[axis] == -1) {
      continue;
    }

    T start = starts[i];
    T end = ends[i];
    T step = steps == nullptr ? 1 : (*steps)[i];

    if (step > 0) {
      slice_dims[axis] = (end - start + step - 1) / step;
    } else {
      slice_dims[axis] = (end - start + step + 1) / step;
    }
  }
  return slice_dims;
}

template <typename T = int64_t>
inline DDim GetDecreasedDims(const DDim slice_dims,
                             const std::vector<T>& decrease_axes,
                             std::vector<T>* infer_flags = nullptr) {
  DDim decreased_dims(slice_dims);
  std::vector<uint8_t> decrease_flag(slice_dims.size(), 0);
  if (decrease_axes.size() > 0) {
    for (size_t i = 0; i < decrease_axes.size(); ++i) {
      T axis = decrease_axes[i];
      decrease_flag[axis] = 1;
      if (infer_flags && (*infer_flags)[i] != -1) {
        PADDLE_ENFORCE_EQ(decreased_dims[axis],
                          1,
                          common::errors::InvalidArgument(
                              "Decrease dim should be 1, but now received %d",
                              decreased_dims[axis]));
      }
    }

    std::vector<T> new_shape;
    for (int i = 0; i < decreased_dims.size(); ++i) {
      if (decrease_flag[i] == 0) {
        new_shape.push_back(decreased_dims[i]);
      }
    }
    decreased_dims = make_ddim(new_shape);
  }
  return decreased_dims;
}

template <typename T = int64_t>
inline void GetDecreasedDimsAndStrides(const std::vector<T> slice_dims,
                                       const std::vector<T> slice_strides,
                                       const std::vector<T>& decrease_axes,
                                       const std::vector<T>& none_axes,
                                       std::vector<T>* new_dims,
                                       std::vector<T>* new_strides,
                                       std::vector<T>* infer_flags = nullptr) {
  std::vector<uint8_t> decrease_flag(slice_dims.size(), 0);
  if (none_axes.size() > 0) {
    size_t none_axes_cur = 0, decrease_axes_cur = 0;
    for (int i = 0; i < slice_dims.size(); ++i) {
      while (none_axes_cur < none_axes.size() &&
             none_axes[none_axes_cur] <= i) {
        new_dims->push_back(1);
        new_strides->push_back(slice_strides[i]);
        none_axes_cur++;
      }
      if (decrease_axes_cur < decrease_axes.size() &&
          decrease_axes[decrease_axes_cur] == i) {
        decrease_axes_cur++;
      } else {
        new_dims->push_back(slice_dims[i]);
        new_strides->push_back(slice_strides[i]);
      }
    }
    while (none_axes_cur < none_axes.size()) {
      new_dims->push_back(1);
      new_strides->push_back(slice_strides[slice_strides.size() - 1]);
      none_axes_cur++;
    }
  } else if (decrease_axes.size() > 0) {
    for (size_t i = 0; i < decrease_axes.size(); ++i) {
      T axis = decrease_axes[i];
      decrease_flag[axis] = 1;
      if (infer_flags && (*infer_flags)[i] != -1) {
        PADDLE_ENFORCE_EQ(slice_dims[axis],
                          1,
                          common::errors::InvalidArgument(
                              "Decrease dim should be 1, but now received %d",
                              slice_dims[axis]));
      }
    }

    for (int i = 0; i < slice_dims.size(); ++i) {
      if (decrease_flag[i] == 0) {
        new_dims->push_back(slice_dims[i]);
        new_strides->push_back(slice_strides[i]);
      }
    }
  } else {
    for (int i = 0; i < slice_dims.size(); ++i) {
      new_dims->push_back(slice_dims[i]);
      new_strides->push_back(slice_strides[i]);
    }
  }

  return;
}

template <typename T = int64_t>
inline void CheckAndUpdateSparseSliceAttrs(const DDim in_dims,
                                           std::vector<T>* axes,
                                           std::vector<T>* starts,
                                           std::vector<T>* ends) {
  int64_t rank = int64_t(in_dims.size());
  for (auto& axis : *axes) {
    if (axis < 0) {
      axis = std::max(int64_t(0), axis + rank);
    }
  }

  PADDLE_ENFORCE_EQ(
      axes->size(),
      starts->size(),
      common::errors::InvalidArgument(
          "The length of axes (%d) and length of starts (%d) should be same.",
          axes->size(),
          starts->size()));
  PADDLE_ENFORCE_EQ(
      axes->size(),
      ends->size(),
      common::errors::InvalidArgument(
          "The length of axes (%d) and length of ends (%d) should be same.",
          axes->size(),
          ends->size()));

  CheckAndUpdateSliceAttrs<T>(in_dims, *axes, starts, ends);
}

inline void ConstructNewSliceAttrs(const DDim& x_dims,
                                   const std::vector<int64_t>& axes,
                                   const std::vector<int64_t>& starts,
                                   const std::vector<int64_t>& ends,
                                   std::vector<int64_t>* new_axes,
                                   std::vector<int64_t>* new_starts,
                                   std::vector<int64_t>* new_ends) {
  for (int64_t i = 0; i < x_dims.size(); ++i) {
    int pos = -1;
    for (int j = 0; j < static_cast<int>(axes.size()); ++j) {
      if (axes[j] == i) {
        pos = j;
        break;
      }
    }
    if (pos == -1) {
      (*new_axes)[i] = i;
      (*new_starts)[i] = 0;
      (*new_ends)[i] = x_dims[i];
    } else {
      (*new_axes)[i] = axes[pos];
      (*new_starts)[i] = starts[pos];
      (*new_ends)[i] = ends[pos];
    }
  }
}

}  // namespace funcs
}  // namespace phi