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

#pragma once
#include <gtest/gtest.h>

#include <mlir/Dialect/Tensor/IR/Tensor.h>
#include <mlir/IR/Builders.h>
#include <mlir/IR/MLIRContext.h>
#include <mlir/IR/Types.h>
#include <mlir/Parser/Parser.h>
#include <mlir/Pass/PassManager.h>

#include "vpux/compiler/dialect/IE/IR/dialect.hpp"
#include "vpux/compiler/dialect/IE/IR/ops/convolution.hpp"
#include "vpux/compiler/dialect/IE/IR/ops/eltwise.hpp"
#include "vpux/compiler/dialect/IE/IR/ops/pooling.hpp"
#include "vpux/compiler/dialect/IE/IR/ops/reduce.hpp"
#include "vpux/compiler/dialect/VPU/IR/dialect.hpp"
#include "vpux/compiler/dialect/VPU/IR/ops/dpu.hpp"
#include "vpux/compiler/dialect/VPU/interfaces/ppe_factory.hpp"
#include "vpux/compiler/dialect/const/dialect.hpp"
#include "vpux/compiler/dialect/const/ops.hpp"
#include "vpux/compiler/init.hpp"

using namespace vpux;

template <typename ConcreteModel, typename MainOpType>
class LayerWithClampOpModel : public IE::LayerWithPostOpInterface::ExternalModel<ConcreteModel, MainOpType> {
public:
    bool isSupportedPostOp(mlir::Operation*, mlir::Operation*, const LogCb&) const {
        return true;
    }

    bool isSupportedClampOp(mlir::Operation*, mlir::Operation*, const LogCb&) const {
        return true;
    }

    void setLayerClampOp(mlir::Operation*, mlir::Operation*) const {
    }
};

template <class MainOpType>
class LayerWithPostOpUsingBiasAndStaticScaleModel final :
        public LayerWithClampOpModel<LayerWithPostOpUsingBiasAndStaticScaleModel<MainOpType>, MainOpType> {
public:
    bool supportsFuseBiasScale(mlir::Operation*) const {
        return true;
    }
};

template <class MainOpType>
class LayerWithPostOpModel final : public LayerWithClampOpModel<LayerWithPostOpModel<MainOpType>, MainOpType> {};

class MLIR_PpeRegistry : public testing::Test {
public:
    MLIR_PpeRegistry() {
        registry = createDialectRegistry();
        registry.addExtension(+[](mlir::MLIRContext* ctx, IE::IEDialect*) {
            IE::ConvolutionOp::attachInterface<LayerWithPostOpUsingBiasAndStaticScaleModel<IE::ConvolutionOp>>(*ctx);
            IE::TransposedConvolutionOp::attachInterface<
                    LayerWithPostOpUsingBiasAndStaticScaleModel<IE::TransposedConvolutionOp>>(*ctx);
            IE::GroupConvolutionOp::attachInterface<
                    LayerWithPostOpUsingBiasAndStaticScaleModel<IE::GroupConvolutionOp>>(*ctx);
            IE::MaxPoolOp::attachInterface<LayerWithPostOpModel<IE::MaxPoolOp>>(*ctx);
            IE::AvgPoolOp::attachInterface<LayerWithPostOpModel<IE::AvgPoolOp>>(*ctx);
            IE::AddOp::attachInterface<LayerWithPostOpModel<IE::AddOp>>(*ctx);
            IE::SubtractOp::attachInterface<LayerWithPostOpModel<IE::SubtractOp>>(*ctx);
            IE::MultiplyOp::attachInterface<LayerWithPostOpModel<IE::MultiplyOp>>(*ctx);
            IE::MatMulOp::attachInterface<LayerWithPostOpModel<IE::MatMulOp>>(*ctx);
        });
    }

protected:
    mlir::DialectRegistry registry;
};

class VPU_PpeUnitBase : public MLIR_PpeRegistry {
protected:
    VPU_PpeUnitBase(std::unique_ptr<VPU::IPpeFactory>&& ppeIfc)
            : MLIR_PpeRegistry(), _ctx(registry), _loc(mlir::UnknownLoc::get(&_ctx)), _ppeIfc(std::move(ppeIfc)) {
        _ctx.loadDialect<Const::ConstDialect>();
        _ctx.loadDialect<IE::IEDialect>();
        _ctx.loadDialect<VPU::VPUDialect>();
    }

protected:
    mlir::MLIRContext _ctx;
    mlir::Location _loc;
    VPU::PpeIfcPtr _ppeIfc;

protected:
    mlir::Type getF32Type() {
        return mlir::Float32Type::get(&_ctx);
    }

    mlir::Type getF16Type() {
        return mlir::Float16Type::get(&_ctx);
    }

    mlir::Type getU8Type() {
        return mlir::quant::UniformQuantizedType::getChecked(_loc, 0, getInt8Type(&_ctx), mlir::Float16Type::get(&_ctx),
                                                             0.002, 128, 0, 255);
    }

    mlir::Type getU8PerAxisType() {
        return mlir::quant::UniformQuantizedPerAxisType::getChecked(
                _loc, 0, getInt8Type(&_ctx), mlir::Float16Type::get(&_ctx), SmallVector<double>({0.002, 0.004}),
                SmallVector<int64_t>({128, 128}), 1, 0, 255);
    }

    mlir::Type getF8Type() {
        return mlir::quant::UniformQuantizedType::getChecked(_loc, 0, mlir::Float8E4M3FNType::get(&_ctx),
                                                             mlir::Float16Type::get(&_ctx), 0.002, 0, -448, 448);
    }

    VPU::PPEStubAttr getStubPPEAttr() {
        return VPU::PPEStubAttr::get(&_ctx);
    }

    template <typename PostOpAttrT, typename... Args>
    PostOpAttrT create(Args&&... args) {
        return PostOpAttrT::get(&_ctx, getFPAttr(&_ctx, std::forward<Args>(args))...);
    }

    Const::DeclareOp createBias(mlir::OpBuilder builder, ArrayRef<type::float16> bias) {
        const auto shape = mlir::RankedTensorType::get({1, static_cast<int64_t>(bias.size()), 1, 1}, getF16Type());
        auto content = Const::ContentAttr::get(Const::createConstContent(shape, bias));

        return builder.create<Const::DeclareOp>(mlir::UnknownLoc::get(&_ctx), content.getType(), std::move(content));
    }

    IE::AddOp createAdd(mlir::OpBuilder builder, mlir::Type in1ElemType, mlir::Type in2ElemType, mlir::Type outElemType,
                        IE::PostOpAttr postOpAttr) {
        auto input1 = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                            ArrayRef<int64_t>{1, 16, 32, 32}, in1ElemType);
        auto input2 = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                            ArrayRef<int64_t>{1, 16, 32, 32}, in2ElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);

        const auto broadcast = IE::AutoBroadcastTypeAttr::get(&_ctx, IE::AutoBroadcastType::NONE_OR_EXPLICIT);

        return builder.create<IE::AddOp>(_loc, outType, input1.getResult(), input2.getResult(), broadcast, postOpAttr,
                                         /*clamp=*/nullptr,
                                         /*outputPadding=*/nullptr,
                                         /*inputPadding=*/nullptr);
    }

    IE::SubtractOp createSubtract(mlir::OpBuilder builder, mlir::Type in1ElemType, mlir::Type in2ElemType,
                                  mlir::Type outElemType, IE::PostOpAttr postOpAttr) {
        auto input1 = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                            ArrayRef<int64_t>{1, 16, 32, 32}, in1ElemType);
        auto input2 = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                            ArrayRef<int64_t>{1, 16, 32, 32}, in2ElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);

        const auto broadcast = IE::AutoBroadcastTypeAttr::get(&_ctx, IE::AutoBroadcastType::NONE_OR_EXPLICIT);

        return builder.create<IE::SubtractOp>(_loc, outType, input1.getResult(), input2.getResult(), broadcast,
                                              postOpAttr, /*clamp=*/nullptr,
                                              /*outputPadding=*/nullptr,
                                              /*inputPadding=*/nullptr);
    }

    IE::MultiplyOp createMultiply(mlir::OpBuilder builder, mlir::Type in1ElemType, mlir::Type in2ElemType,
                                  mlir::Type outElemType, IE::PostOpAttr postOpAttr) {
        auto input1 = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                            ArrayRef<int64_t>{1, 16, 32, 32}, in1ElemType);
        auto input2 = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                            ArrayRef<int64_t>{1, 16, 32, 32}, in2ElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);

        const auto broadcast = IE::AutoBroadcastTypeAttr::get(&_ctx, IE::AutoBroadcastType::NONE_OR_EXPLICIT);

        return builder.create<IE::MultiplyOp>(_loc, outType, input1.getResult(), input2.getResult(), broadcast,
                                              postOpAttr, /*clamp=*/nullptr,
                                              /*outputPadding=*/nullptr,
                                              /*inputPadding=*/nullptr);
    }

    IE::ConvolutionOp createConvolution(mlir::OpBuilder builder, mlir::Type inElemType, mlir::Type weightsElemType,
                                        mlir::Type outElemType, double scale, IE::PostOpAttr postOpAttr,
                                        Const::DeclareOp bias) {
        auto input = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                           ArrayRef<int64_t>{1, 16, 32, 32}, inElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);
        auto weights = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                             ArrayRef<int64_t>{1, 16, 32, 32}, weightsElemType);

        const auto strides = getIntArrayAttr(&_ctx, SmallVector<int64_t>{1, 1});
        const auto padsBegin = getIntArrayAttr(&_ctx, SmallVector<int64_t>{0, 0});
        const auto padsEnd = getIntArrayAttr(&_ctx, SmallVector<int64_t>{0, 0});
        const auto dilations = getIntArrayAttr(&_ctx, SmallVector<int64_t>{1, 1});
        const auto staticScale = getFPAttr(&_ctx, scale);

        return builder.create<IE::ConvolutionOp>(
                _loc, outType, input.getResult(), weights.getResult(), bias != nullptr ? bias.getResult() : nullptr,
                /*scale*/ nullptr, strides, padsBegin, padsEnd, dilations, postOpAttr, /*clamp=*/nullptr, staticScale,
                /*outputPadding=*/nullptr, /*inputPadding=*/nullptr);
    }

    IE::AvgPoolOp createAvgPool(mlir::OpBuilder builder, mlir::Type inElemType, mlir::Type outElemType,
                                ArrayRef<int64_t> kernelShape, double scale, IE::PostOpAttr postOpAttr) {
        auto input = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                           ArrayRef<int64_t>{1, 16, 32, 32}, inElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);

        const auto strides = getIntArrayAttr(&_ctx, SmallVector<int64_t>{1, 1});
        const auto kernel = getIntArrayAttr(&_ctx, kernelShape);
        const auto padsBegin = getIntArrayAttr(&_ctx, SmallVector<int64_t>{0, 0});
        const auto padsEnd = getIntArrayAttr(&_ctx, SmallVector<int64_t>{0, 0});
        const auto rounding = IE::RoundingTypeAttr::get(&_ctx, IE::RoundingType::FLOOR);
        const auto staticScale = getFPAttr(&_ctx, scale);

        return builder.create<IE::AvgPoolOp>(_loc, outType, input.getResult(), kernel, strides, padsBegin, padsEnd,
                                             rounding, nullptr, postOpAttr, /*clamp=*/nullptr, staticScale,
                                             /*outputPadding=*/nullptr, /*inputPadding=*/nullptr);
    }

    IE::MaxPoolOp createMaxPool(mlir::OpBuilder builder, mlir::Type inElemType, mlir::Type outElemType,
                                IE::PostOpAttr postOpAttr) {
        auto input = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                           ArrayRef<int64_t>{1, 16, 32, 32}, inElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);

        const auto strides = getIntArrayAttr(&_ctx, SmallVector<int64_t>{1, 1});
        const auto kernel = getIntArrayAttr(&_ctx, SmallVector<int64_t>{2, 2});
        const auto padsBegin = getIntArrayAttr(&_ctx, SmallVector<int64_t>{0, 0});
        const auto padsEnd = getIntArrayAttr(&_ctx, SmallVector<int64_t>{0, 0});
        const auto rounding = IE::RoundingTypeAttr::get(&_ctx, IE::RoundingType::FLOOR);

        return builder.create<IE::MaxPoolOp>(_loc, outType, input.getResult(), kernel, strides, padsBegin, padsEnd,
                                             rounding, postOpAttr, /*clamp=*/nullptr,
                                             /*outputPadding=*/nullptr,
                                             /*inputPadding=*/nullptr);
    }

    IE::MatMulOp createMatMul(mlir::OpBuilder builder, mlir::Type in1ElemType, mlir::Type in2ElemType,
                              mlir::Type outElemType) {
        auto input1 = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                            ArrayRef<int64_t>{1, 16, 32, 32}, in1ElemType);
        auto input2 = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                            ArrayRef<int64_t>{1, 16, 32, 32}, in2ElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);

        return builder.create<IE::MatMulOp>(_loc, outType, input1.getResult(), input2.getResult(), false, false);
    }

    IE::ReduceMeanOp createReduceMean(mlir::OpBuilder builder, mlir::Type inElemType, mlir::Type outElemType,
                                      ArrayRef<int64_t> axes, ArrayRef<int64_t> inputPadding) {
        auto input = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                           ArrayRef<int64_t>{1, 16, 32, 32}, inElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);
        const auto axesAttr = getIntArrayAttr(&_ctx, axes);
        const auto inputPaddingAttr = inputPadding.empty() ? nullptr : getIntArrayAttr(&_ctx, inputPadding);

        return builder.create<IE::ReduceMeanOp>(_loc, outType, input.getResult(), /*axes=*/nullptr, axesAttr,
                                                /*keep_dims=*/nullptr, /*output_padding=*/nullptr, inputPaddingAttr);
    }

    VPU::NCEReduceOp createNCEReduce(mlir::OpBuilder builder, mlir::Type inElemType, mlir::Type outElemType,
                                     ArrayRef<int64_t> axes, VPU::ReduceType reduceType,
                                     ArrayRef<int64_t> inputPadding) {
        auto input = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                           ArrayRef<int64_t>{1, 16, 32, 32}, inElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);
        const auto axesAttr = getIntArrayAttr(&_ctx, axes);
        const auto typeAttr = VPU::ReduceTypeAttr::get(&_ctx, reduceType);
        const auto inputPaddingAttr = inputPadding.empty() ? nullptr : getIntArrayAttr(&_ctx, inputPadding);
        return builder.create<VPU::NCEReduceOp>(_loc, outType, input.getResult(), axesAttr, /*ppe=*/nullptr, typeAttr,
                                                /*multiClusterStrategy=*/nullptr, /*output_padding=*/nullptr,
                                                inputPaddingAttr);
    }

    VPU::NCEInterpolateOp createNCEInterpolate(mlir::OpBuilder builder, mlir::Type inElemType,
                                               mlir::Type weightsElemType, mlir::Type outElemType,
                                               VPU::PPEAttr oldPpeAttr) {
        auto input = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                           ArrayRef<int64_t>{1, 16, 32, 32}, inElemType);
        auto weights = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                             ArrayRef<int64_t>{1, 16, 32, 32}, weightsElemType);
        auto weightsTable = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                                  ArrayRef<int64_t>{16, 1, 1, 4}, weightsElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);

        const auto strides = getIntArrayAttr(&_ctx, SmallVector<int64_t>{1, 1});
        const auto rawFilterShape = getIntArrayAttr(&_ctx, mlir::cast<NDTypeInterface>(weights.getType()).getShape());

        return builder.create<VPU::NCEInterpolateOp>(_loc, outType, input, weights, weightsTable, nullptr, nullptr,
                                                     nullptr, nullptr, nullptr, strides, oldPpeAttr, rawFilterShape,
                                                     /*multiClusterStrategy=*/nullptr,
                                                     /*mode=*/nullptr);
    }

    VPU::NCEDepthConvolutionOp createNCEDWConv(mlir::OpBuilder builder, mlir::Type inElemType,
                                               mlir::Type weightsElemType, mlir::Type outElemType,
                                               VPU::PPEAttr oldPpeAttr) {
        auto input = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                           ArrayRef<int64_t>{1, 16, 32, 32}, inElemType);
        auto weights = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                             ArrayRef<int64_t>{16, 16, 1, 1}, weightsElemType);
        auto weightsTable = builder.create<mlir::tensor::EmptyOp>(mlir::UnknownLoc::get(&_ctx),
                                                                  ArrayRef<int64_t>{16, 1, 1, 4}, weightsElemType);
        const auto outType = mlir::RankedTensorType::get(ArrayRef<int64_t>{1, 16, 32, 32}, outElemType);

        const auto strides = getIntArrayAttr(&_ctx, SmallVector<int64_t>{1, 1});
        const auto rawFilterShape = getIntArrayAttr(&_ctx, SmallVector<int64_t>{16, 1, 3, 3});

        const auto padsBegin = getIntArrayAttr(&_ctx, SmallVector<int64_t>{1, 1});
        const auto padsEnd = getIntArrayAttr(&_ctx, SmallVector<int64_t>{1, 1});
        const auto padAttr = VPU::getPaddingAttr(&_ctx, PadInfo(padsBegin, padsEnd));

        return builder.create<VPU::NCEDepthConvolutionOp>(
                _loc, outType, input, weights, weightsTable, /*dataPointerTensor=*/nullptr,
                /*sparsityPointerTensor=*/nullptr,
                /*scaleTensor=*/nullptr, /*biasTensor=*/nullptr, /*zeroPointTensor=*/nullptr, strides, padAttr,
                oldPpeAttr, rawFilterShape,
                /*multiClusterStrategy=*/nullptr,
                /*outputPaddingAttr=*/nullptr,
                /*inputPaddingAttr=*/nullptr);
    }
};
