/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include "opdev/make_op_executor.h"
#include "opdev/op_dfx.h"
#include "opdev/tensor_view_utils.h"

namespace l0op {

OP_TYPE_REGISTER(SqueezeNd);

bool SqueezeNdInferShape(const op::Shape &shape, const aclIntArray* dim, op::Shape &newShape) {
  auto dimNum = static_cast<int64_t>(shape.GetDimNum());
  if (dim->Size() == 0) {
    for (auto i = 0; i < dimNum; i++) {
      if (shape[i] != 1) {
        newShape.AppendDim(shape[i]);
      }
    }
    return true;
  }
  op::FVector<bool, op::MAX_DIM_NUM> squeezeDim(dimNum, false);
  for (uint64_t i = 0; i < dim->Size(); i++) {
    auto idx = (*dim)[i];
    idx = idx >= 0 ? idx : dimNum + idx;
    if (idx >= dimNum || idx < 0) {
      OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Invalid dim value for squeeze operation. shape: %s, dims: %s",
              op::ToString(shape).GetString(),
              dim->ToString().GetString());
      return false;
    }
    squeezeDim[idx] = shape[idx] == 1;
  }

  for (auto i = 0; i < dimNum; i++) {
    if (!squeezeDim[i]) {
      newShape.AppendDim(shape[i]);
    }
  }
  return true;
}

const aclTensor *SqueezeNd(const aclTensor *x, const aclIntArray* dim, aclOpExecutor *executor) {
  L0_DFX(SqueezeNd, x, dim);
  if (!op::IsContiguous(x)) {
    OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Input tensor must be contiguous.");
    return nullptr;
  }
  op::Shape newShape;
  if (!SqueezeNdInferShape(x->GetViewShape(), dim, newShape)) {
    OP_LOGE(ACLNN_ERR_PARAM_INVALID, "InferShape failed for squeeze operation.");
    return nullptr;
  }
  return executor->CreateView(x, newShape, x->GetViewOffset());
}

const aclTensor *SqueezeNd(const aclTensor *x, int64_t dim, aclOpExecutor *executor) {
  return SqueezeNd(x, executor->AllocIntArray(&dim, 1), executor);
}
} // namespace l0op