//
// Copyright (C) 2023-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#pragma once

#include "vpux/compiler/dialect/IE/IR/attributes.hpp"
#include "vpux/compiler/dialect/const/utils/content.hpp"

namespace vpux {
namespace IE {

SmallVector<int64_t> getBroadcastAxesNumpyBidirectional(ArrayRef<int64_t> inputShape, ArrayRef<int64_t> outputShape);
SmallVector<int64_t> getBroadcastAxesExplicit(ArrayRef<int64_t> axesMapping, ArrayRef<int64_t> outputShape);

mlir::Value createShapeConstForBroadcast(mlir::OpBuilder& rewriter, mlir::Location loc, ShapeRef shape);
mlir::Value createBroadcast(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value input, ShapeRef targetShape,
                            mlir::Value axisMapping = nullptr, IE::BroadcastTypeAttr broadcastTypeAttr = nullptr);
mlir::Value createDynamicBroadcast(mlir::OpBuilder& rewriter, mlir::Location loc, mlir::Value input, mlir::Value other,
                                   mlir::Value axisMapping = nullptr,
                                   IE::BroadcastTypeAttr broadcastTypeAttr = nullptr);

/// Aligns shapes of two buffers by broadcasting across misaligned axes in order
/// to allow 'X op Y' operation to behave correctly. This procedure could fail
/// when x's and y's shapes differ and both are non-1 at any given axis (this
/// situation is ambiguous).
///
/// Note: this procedure does not guarantee any state consistency if it fails
/// "mid-way", that is, any operands (x or y) could still be modified (e.g.
/// broadcasted across some axes).
mlir::LogicalResult broadcastAlignShapes(mlir::MLIRContext* ctx, Const::Content& x, Const::Content& y,
                                         const Logger& log);
}  // namespace IE
}  // namespace vpux
