/* ******************************************************************************
 *
 *
 * This program and the accompanying materials are made available under the
 * terms of the Apache License, Version 2.0 which is available at
 * https://www.apache.org/licenses/LICENSE-2.0.
 *
 *  See the NOTICE file distributed with this work for additional
 *  information regarding copyright ownership.
 * 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.
 *
 * SPDX-License-Identifier: Apache-2.0
 ******************************************************************************/

//
// @author Yurii Shyrma (iuriish@yahoo.com), created on 16.04.2019
//

#ifndef LIBND4J_LOOPKIND_H
#define LIBND4J_LOOPKIND_H

#include <helpers/shape.h>


namespace sd {

class SD_LIB_EXPORT LoopKind {
 public:
  enum Kind {
    RANK1,
    RANK2,
    RANK3,
    RANK4,
    RANK5,
    BROADCAST_SCALAR_X,
    BROADCAST_SCALAR_Y,
    BROADCAST_2D,    // Added new 2D broadcast case
    BROADCAST_3D,
    BROADCAST_4D,
    BROADCAST_5D,
    COMMON,
    SMALLARR2DX
  };


  static SD_INLINE Kind deduceKindOfLoopXZ(const LongType* xShapeInfo, const LongType* zShapeInfo);
  static SD_INLINE Kind deduceKindOfLoopXYZ(const LongType* xShapeInfo, const LongType* yShapeInfo,
                                            const LongType* zShapeInfo);
  static SD_INLINE Kind deduceKindOfLoopTadXZ(const LongType* xShapeInfo, const LongType* zShapeInfo,
                                              const LongType* tadShapeInfo);
  static SD_INLINE Kind deduceKindOfLoopTadXYZ(const LongType* xTadShapeInfo, const LongType* yTadShapeInfo,
                                               const LongType* zShapeInfo);
  static SD_INLINE Kind deduceKindOfLoopBroadcast(const LongType* xShapeInfo, const LongType* yShapeInfo,
                                                  const LongType* zShapeInfo);
};

//////////////////////////////////////////////////////////////////////////////
LoopKind::Kind LoopKind::deduceKindOfLoopXZ(const LongType* xShapeInfo, const LongType* zShapeInfo) {
  const int xRank = shape::rank(xShapeInfo);
  const bool shapesSame = shape::shapeEquals(xShapeInfo, zShapeInfo);

  // Handle rank-specific optimizations when shapes match
  if (shapesSame) {
    switch(xRank) {
      case 1: return RANK1;
      case 2: return RANK2;
      case 3: return RANK3;
      case 4: return RANK4;
      case 5: return RANK5;
      default: return COMMON;
    }
  }

  return COMMON;
}

LoopKind::Kind LoopKind::deduceKindOfLoopBroadcast(const LongType* xShapeInfo, const LongType* yShapeInfo,
                                                   const LongType* zShapeInfo) {
  auto xRank = shape::rank(xShapeInfo);
  auto yRank = shape::rank(yShapeInfo);
  auto zRank = shape::rank(zShapeInfo);

  auto xOrder = shape::order(xShapeInfo);
  auto yOrder = shape::order(yShapeInfo);
  auto zOrder = shape::order(zShapeInfo);

  // First check scalar broadcast cases
  if (yRank < 1 && xRank == yRank && xRank == zRank && xOrder == 'c' && yOrder == 'c' && zOrder == 'c' && xRank >= 2) {
    // Validate shapes are equal till last dim
    for (int e = 0; e < xRank - 1; e++) {
      if (xShapeInfo[e + 1] != yShapeInfo[e + 1]) break;
    }

    // Check if one shape has 1 as last dim
    auto detect = xShapeInfo[xRank] == 1 ? -1 : (yShapeInfo[xRank] == 1) ? 1 : 0;

    if (detect == 1) return BROADCAST_SCALAR_Y;
    else if (detect == -1) return BROADCAST_SCALAR_X;
  }

  // Check for 2D broadcasting cases
  if (zRank == 2) {
    const auto zShape = shape::shapeOf(zShapeInfo);

    // Case 1: Matrix + row vector broadcasting
    if ((xRank == 2 && yRank == 1) || (yRank == 2 && xRank == 1)) {
      const auto vecShape = xRank == 1 ? shape::shapeOf(xShapeInfo) : shape::shapeOf(yShapeInfo);
      if (vecShape[0] == zShape[1]) return BROADCAST_2D;
    }

    // Case 2: Matrix + column vector broadcasting (nx1)
    if (xRank == 2 && yRank == 2) {
      const auto xShape = shape::shapeOf(xShapeInfo);
      const auto yShape = shape::shapeOf(yShapeInfo);

      // Improved check for column vector - explicitly check for yShape[1] == 1
      if (yShape[0] == zShape[0] && yShape[1] == 1) {
        return BROADCAST_2D;
      }

      // Also check for row vector - explicitly check for yShape[0] == 1
      if (yShape[0] == 1 && yShape[1] == zShape[1]) {
        return BROADCAST_2D;
      }
    }

    // Case 3: Regular 2D broadcasting with matching ranks
    if (xRank == 2 && yRank == 2) {
      const auto xShape = shape::shapeOf(xShapeInfo);
      const auto yShape = shape::shapeOf(yShapeInfo);

      // Check if one array can broadcast to the other
      bool canBroadcast = (xShape[0] == zShape[0] || xShape[0] == 1) &&
                          (xShape[1] == zShape[1] || xShape[1] == 1) &&
                          (yShape[0] == zShape[0] || yShape[0] == 1) &&
                          (yShape[1] == zShape[1] || yShape[1] == 1);

      if (canBroadcast) return BROADCAST_2D;
    }
  }

  // Check higher dimension cases
  bool bNDLoopsRanks = (xRank == zRank && yRank <= xRank && yRank >= 2);

  int countUnityDimsInY = 0, countUnityDimsInX = 0;
  for (LongType i = 0; i < xRank; i++) {
    if (i < yRank) countUnityDimsInY += (1 == shape::sizeAt(yShapeInfo, i)) ? 1 : 0;
    countUnityDimsInX += (1 == shape::sizeAt(xShapeInfo, i)) ? 1 : 0;
  }


  if (3 == xRank) return BROADCAST_3D;
  if (4 == xRank) return BROADCAST_4D;
  if (5 == xRank) return BROADCAST_5D;


  return COMMON;
}
//////////////////////////////////////////////////////////////////////////////
LoopKind::Kind LoopKind::deduceKindOfLoopXYZ(const LongType* xShapeInfo, const LongType* yShapeInfo,
                                             const LongType* zShapeInfo) {
  const int xRank = shape::rank(xShapeInfo);
  const char xOrder = shape::order(xShapeInfo);
  const char yOrder = shape::order(yShapeInfo);
  const char zOrder = shape::order(zShapeInfo);

  // Check if all shapes match
  const bool shapesSame = shape::shapeEquals(xShapeInfo, yShapeInfo, zShapeInfo);

  // Handle rank-specific optimizations when shapes match
  if (shapesSame) {
    if (xRank == 1) return RANK1;
    if (xRank == 2) return RANK2;
    if (xRank == 3) return RANK3;
    if (xRank == 4) return RANK4;
    if (xRank == 5) return RANK5;
  }

  // Default case
  return COMMON;
}

//////////////////////////////////////////////////////////////////////////////
LoopKind::Kind LoopKind::deduceKindOfLoopTadXZ(const LongType* xShapeInfo, const LongType* zShapeInfo,
                                               const LongType* tadShapeInfo) {
  // Check for small array optimization first
  if (shape::rank(xShapeInfo) == 2 &&
      shape::length(tadShapeInfo) * shape::length(zShapeInfo) <= Environment::getInstance().elementwiseThreshold()) {
    return SMALLARR2DX;
  }

  // Handle rank-specific optimizations
  switch(shape::rank(tadShapeInfo)) {
    case 1: return RANK1;
    case 2: return RANK2;
    case 3: return RANK3;
    case 4: return RANK4;
    case 5: return RANK5;
    default: return COMMON;
  }
}
//////////////////////////////////////////////////////////////////////////////
LoopKind::Kind LoopKind::deduceKindOfLoopTadXYZ(const LongType* xTadShapeInfo, const LongType* yTadShapeInfo,
                                                const LongType* zShapeInfo) {
  // both tad shapes are the same, but strides may be different
  const int tadRank = shape::rank(xTadShapeInfo);

  // Handle rank-specific optimizations
  switch(tadRank) {
    case 1: return RANK1;
    case 2: return RANK2;
    case 3: return RANK3;
    case 4: return RANK4;
    case 5: return RANK5;
    default: return COMMON;
  }
}

}  // namespace sd
#endif  // LIBND4J_LOOPKIND_H