// Copyright Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#include <rocRoller/AssemblyKernel.hpp>
#include <rocRoller/Context.hpp>
#include <rocRoller/Expression.hpp>
#include <rocRoller/ExpressionTransformations.hpp>
#include <rocRoller/KernelOptions_detail.hpp>
#include <rocRoller/Operations/Command.hpp>

#include "CustomMatchers.hpp"
#include "CustomSections.hpp"
#include "ExpressionMatchers.hpp"
#include "TestContext.hpp"
#include "TestKernels.hpp"

#include <common/TestValues.hpp>

#include <catch2/catch_test_macros.hpp>

namespace FastDivisionTest
{
    auto getMagicMultiple(auto x)
    {
        using namespace rocRoller;
        namespace Ex = Expression;
        auto expr    = magicMultiple(Ex::literal(x));

        return Ex::literal(evaluate(expr));
    }

    TEST_CASE("FastDivision ExpressionTransformation works for constant expressions.",
              "[division][expression][expression-transformation]")
    {
        using namespace rocRoller;
        namespace Ex = Expression;

        auto context = TestContext::ForDefaultTarget({{.minLaunchTimeExpressionComplexity = 49}});

        auto command = std::make_shared<Command>();

        auto aTag = command->allocateTag();
        auto a    = command
                     ->allocateArgument(
                         {DataType::Int32, PointerType::Value}, aTag, ArgumentType::Value)
                     ->expression();

        auto b = Register::Value::Placeholder(
                     context.get(), Register::Type::Vector, DataType::Int32, 1)
                     ->expression();

        auto argsBefore = context->kernel()->arguments().size();

        auto expr      = a / Ex::literal(8u);
        auto expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast, IdenticalTo(a >> Ex::literal(3u)));

        expr      = a / Ex::literal(8);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast,
                   IdenticalTo((a + Ex::logicalShiftR((a >> Ex::literal(31u)), Ex::literal(29u))
                                >> Ex::literal(3))));

        expr      = b / Ex::literal(7u);
        expr_fast = fastDivision(expr, context.get());

        {
            auto mulHigh = multiplyHigh(b, getMagicMultiple(7u));
            CHECK_THAT(
                expr_fast,
                EquivalentTo((((b - mulHigh) >> Ex::literal(1u)) + mulHigh) >> Ex::literal(2)));
        }

        expr      = b / Ex::literal(1);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(b));

        expr      = b / Ex::literal(-5);
        expr_fast = fastDivision(expr, context.get());

        {
            auto mulPlusB = multiplyHigh(b, getMagicMultiple(-5)) + b;
            CHECK_THAT(expr_fast,
                       EquivalentTo((((mulPlusB + ((mulPlusB >> Ex::literal(31)) & Ex::literal(4)))
                                      >> Ex::literal(2u))
                                     ^ Ex::literal(-1))
                                    + Ex::literal(1)));
        }

        expr      = a / Ex::literal(8u);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(a >> Ex::literal(3u)));

        expr      = a / Ex::literal(128u);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(a >> Ex::literal(7u)));

        expr      = a / (Ex::literal(43u) + Ex::literal(85u));
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(a >> Ex::literal(7u)));

        auto argsAfter = context->kernel()->arguments().size();
        CHECK(argsBefore == argsAfter);
    }

    TEST_CASE("FastDivision ExpressionTransformation works for argument expressions.",
              "[division][expression][expression-transformation]")
    {
        using namespace rocRoller;
        namespace Ex = Expression;

        auto context = TestContext::ForDefaultTarget({{.minLaunchTimeExpressionComplexity = 49}});

        auto command = std::make_shared<rocRoller::Command>();

        auto reg = Register::Value::Placeholder(
            context.get(), Register::Type::Vector, DataType::Int32, 1);
        auto a = reg->expression();

        auto reg2 = Register::Value::Placeholder(
            context.get(), Register::Type::Vector, DataType::UInt32, 1);
        auto a_unsigned = reg2->expression();

        auto bSignedTag = command->allocateTag();
        auto b_signed   = command
                            ->allocateArgument({DataType::Int32, PointerType::Value},
                                               bSignedTag,
                                               ArgumentType::Value)
                            ->expression();

        auto bUnsignedTag = command->allocateTag();
        auto b_unsigned   = command
                              ->allocateArgument({DataType::UInt32, PointerType::Value},
                                                 bUnsignedTag,
                                                 ArgumentType::Value)
                              ->expression();

        auto expr      = a / b_signed;
        auto expr_fast = fastDivision(expr, context.get());

        {
            auto [mul, shift, sign, _] = getMagicDivisionParams(b_signed, context.get());

            auto mulPlusA = a + multiplyHigh(a, mul);

            CHECK_THAT(expr_fast,
                       EquivalentTo(
                           ((((mulPlusA
                               + ((mulPlusA >> Ex::literal(31))
                                  & ((Ex::literal(1) << shift)
                                     + Ex::conditional(
                                         mul == Ex::literal(0), Ex::literal(-1), Ex::literal(0)))))
                              >> shift)
                             ^ sign)
                            - sign)));
        }

        expr      = a_unsigned / b_unsigned;
        expr_fast = fastDivision(expr, context.get());

        {
            auto [mul, shift, sign, shiftMSB] = getMagicDivisionParams(b_unsigned, context.get());

            auto mulHigh = multiplyHigh(a_unsigned, mul);

            auto t = ((a_unsigned - mulHigh >> Ex::literal(1u)) + mulHigh) >> shift;

            auto result = Ex::conditional(shiftMSB == Ex::literal(0u), t, a_unsigned);

            CHECK_THAT(expr_fast, EquivalentTo(result));
        }
    }

    TEST_CASE("FastDivision ExpressionTransformation works for constant modulo expressions.",
              "[division][expression][expression-transformation]")
    {
        using namespace rocRoller;
        namespace Ex = Expression;

        auto context = TestContext::ForDefaultTarget({{.minLaunchTimeExpressionComplexity = 49}});

        auto command = std::make_shared<Command>();

        auto regPtr = Register::Value::Placeholder(
            context.get(), Register::Type::Vector, DataType::UInt32, 1);
        auto reg = regPtr->expression();

        auto aTag = command->allocateTag();
        auto a    = command
                     ->allocateArgument(
                         {DataType::Int32, PointerType::Value}, aTag, ArgumentType::Value)
                     ->expression();

        a = context->kernel()->addArgument({"arg", DataType::Int32, DataDirection::ReadOnly, a});

        auto argsBefore = context->kernel()->arguments();

        auto expr      = a % Ex::literal(8u);
        auto expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(a & Ex::literal(7u)));

        expr      = a % Ex::literal(8);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast,
                   EquivalentTo(a
                                - ((a + logicalShiftR((a >> Ex::literal(31u)), Ex::literal(29u)))
                                   & Ex::literal(-8))));

        expr      = a % Ex::literal(7u);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(
            expr_fast,
            EquivalentTo(a - (fastDivision(a / Ex::literal(7u), context.get()) * Ex::literal(7u))));

        expr      = a % Ex::literal(1);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(Ex::literal(0)));

        expr      = a % Ex::literal(1u);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(Ex::literal(0u)));

        expr      = a % Ex::literal(-5);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(
            expr_fast,
            EquivalentTo(a - (fastDivision(a / Ex::literal(-5), context.get()) * Ex::literal(-5))));

        expr      = a % Ex::literal(8u);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(a & Ex::literal(7u)));

        expr      = a % Ex::literal(128u);
        expr_fast = fastDivision(expr, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(a & Ex::literal(127u)));

        auto argsAfter = context->kernel()->arguments();
        CHECK(argsBefore == argsAfter);
    }

    TEST_CASE("FastDivision ExpressionTransformation works for modulo of argument expressions.",
              "[division][expression][expression-transformation]")
    {
        using namespace rocRoller;
        namespace Ex = Expression;

        auto context = TestContext::ForDefaultTarget({{.minLaunchTimeExpressionComplexity = 49}});

        auto command = std::make_shared<Command>();

        auto reg = Register::Value::Placeholder(
            context.get(), Register::Type::Vector, DataType::Int32, 1);
        auto a = reg->expression();

        auto reg2 = Register::Value::Placeholder(
            context.get(), Register::Type::Vector, DataType::UInt32, 1);
        auto a_unsigned = reg2->expression();

        auto bSignedTag = command->allocateTag();
        auto b_signed   = command
                            ->allocateArgument({DataType::Int32, PointerType::Value},
                                               bSignedTag,
                                               ArgumentType::Value)
                            ->expression();

        auto bUnsignedTag = command->allocateTag();
        auto b_unsigned   = command
                              ->allocateArgument({DataType::UInt32, PointerType::Value},
                                                 bUnsignedTag,
                                                 ArgumentType::Value)
                              ->expression();

        auto expr      = a % b_signed;
        auto expr_fast = fastDivision(expr, context.get());
        auto fastDiv   = fastDivision(a / b_signed, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(a - (fastDiv * b_signed)));

        expr      = a_unsigned % b_unsigned;
        expr_fast = fastDivision(expr, context.get());
        fastDiv   = fastDivision(a_unsigned / b_unsigned, context.get());
        CHECK_THAT(expr_fast, EquivalentTo(a_unsigned - (fastDiv * b_unsigned)));
    }
}