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

#pragma once

#include "vpux/compiler/dialect/VPU/interfaces/sparsity_constraint.hpp"
#include "vpux/compiler/utils/options.hpp"
#include "vpux/utils/core/array_ref.hpp"
#include "vpux/utils/core/small_vector.hpp"

#include <mlir/IR/Operation.h>
#include <mlir/IR/Types.h>

#include <cstdint>
#include <string>

namespace vpux::config {
enum class ArchKind : uint64_t;
}

namespace vpux {
namespace VPU {

enum class EnableActivationSparsityMode { AUTO, TRUE, FALSE };

EnableActivationSparsityMode getActSparsityMode(std::string enableActivationSparsityOption);
EnableActivationSparsityMode getActSparsityMode(const StrOption& enableActivationSparsityOption);
bool isActSparsityEnabled(const StrOption& enableActivationSparsityOption);

int64_t getSESize(int64_t channels, const VPU::SparsityConstraint& sparsityConstraint, bool isDepthwise = false);

/*
    Effective sparse output type is the actual tensor IDU sees at its input after applying SETable over the data.

    For example, for a SEAttr with Interpolate Nearest with 2x2 scales, we'll have the following shapes for the sparse
    type components:
    data: [1, 16, 32, 32]
    sparsity_map: [1, 16, 64, 64]
    storage_element_table [1, 1, 64, 64]

    Effective output type will have shape: [1, 16, 64, 64]
*/
mlir::Type getEffectiveSparseOutputType(mlir::Type sparseType);

enum SparsityRemovalFlag {
    Success,
    ClusteredOpInterfaceMissingFail,
    MultiClusterStrategyMissingFail,
    SOKMissingFail,
    SparseOutputMissingFail,
    CatchAllFail
};

SparsityRemovalFlag shouldRemoveOutputSparsity(mlir::Operation* op);

// VPU5 adds a new feature to support SEP without sparsity bitmap
bool isSEOnlyWithoutSMSupported(config::ArchKind arch);

std::pair<SmallVector<int64_t>, SmallVector<int64_t>> getUpdatedSliceOffsetsAndShapesForSETable(
        int64_t seDepth, mlir::ArrayAttr seSizeAttr, ArrayRef<int64_t> sliceOffsets, ArrayRef<int64_t> sliceSizes);

/*
 * Check if any of the activation inputs or outputs is sparse tensor type
 */
bool isActSparseOp(mlir::Operation* op);

}  // namespace VPU
}  // namespace vpux
