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

#include <rocRoller/Expression.hpp>
#include <rocRoller/KernelGraph/KernelGraph.hpp>
#include <rocRoller/KernelGraph/Transforms/AddPRNG.hpp>
#include <rocRoller/KernelGraph/Utils.hpp>

namespace rocRoller
{
    namespace KernelGraph
    {
        using namespace ControlGraph;
        using namespace CoordinateGraph;

        /**
         *  Get an unordered map where the values are random number generators(PRNG) and
         *  keys are the expressions (Assign) that contain random numbers generated from
         *  the corresponding generators.
         *
         *  Note: an assumption is that an expression (Assign) can only have at most
         *        one random number from each generator.
         */
        auto getMappingOfRandomNumbersAndPRNGs(KernelGraph& kgraph)
        {
            auto& control = kgraph.control;
            auto& mapper  = kgraph.mapper;

            std::unordered_map<int, int> VGPR2Generators; //  SeedVGPR  <--->  Generator
            for(auto prng : control.getNodes<SeedPRNG>())
            {
                auto seedVGPR
                    = mapper.get(prng, NaryArgument::DEST); // seedVGPR stores the seed value
                VGPR2Generators[seedVGPR] = prng;
            }

            std::unordered_map<int, std::vector<int>>
                assign2Generators; // Assign (RandomNumbers) <---> Generators

            if(VGPR2Generators.empty())
            {
                // If the control graph has no PRNG, this transformation is not
                // needed and we can return early to save some time.
                return assign2Generators;
            }

            for(auto assignOp : control.getNodes<Assign>())
            {
                auto elem = control.getNode<Assign>(assignOp);
                for(auto [vgpr, prng] : VGPR2Generators)
                {
                    auto subExpr
                        = Expression::DataFlowTag{vgpr, Register::Type::Vector, DataType::UInt32};
                    auto& expr = *elem.expression;
                    // Search within this expression to see if it contains any random number generated by
                    // this PRNG
                    if(containsSubExpression(expr, subExpr))
                    {
                        assign2Generators[assignOp].push_back(prng);
                    }
                }
            }
            return assign2Generators;
        }

        /** This function creates a sub-coordinate graph which is evaluated to
         *  this expression:
         *
         *  TID =
         *    (workgroupSize_X * workgroupId_X + workItemId_X) +
         *    (workgroupSize_Y * workgroupId_Y + workItemId_Y) +
         *    (workgroupSize_Z * workgroupId_Z + workItemId_Z)
         *
         *
         *   Visualization of the sub-coordinate graph (Z dim is omitted):
         *
         *    Workgroup(x)    Workitem(x)      Workgroup(y)    Workitem(y)
         *         |             |                   |             |
         *         |_____________|                   |_____________|
         *                |                                 |
         *            <Flatten>                         <Flatten>
         *                |                                 |
         *                |                                 |
         *                v                                 v
         *             subdim(x)                        subdim(y)
         *                |                                 |
         *                |_________________________________|
         *                                |
         *                             <Join>
         *                                |
         *                                |
         *                                v
         *                             tidVGPR
         */
        int createTIDSubgraph(ContextPtr                                                context,
                              rocRoller::KernelGraph::CoordinateGraph::CoordinateGraph& coordinates)
        {
            auto        tidVGPR       = coordinates.addElement(VGPR()); // a new VGPR to store TID
            auto const  numDimensions = context->kernel()->kernelDimensions();
            auto const& ws            = context->kernel()->workgroupSize();

            std::vector<int> subdims;
            for(size_t i = 0; i < numDimensions; i++)
            {
                auto wg_tag = coordinates.addElement(Workgroup(i));
                auto wi_tag = coordinates.addElement(Workitem(i, Expression::literal(ws[i])));

                auto subdim
                    = coordinates.addElement(SubDimension(i, nullptr, Expression::literal(1u)));

                coordinates.addElement(Flatten(), {wg_tag, wi_tag}, {subdim});

                subdims.push_back(subdim);
            }

            coordinates.addElement(Join(), subdims, std::vector<int>{tidVGPR});
            return tidVGPR;
        }

        /** Generate and store a random number into the random number generator's
         *  seed VPGR before each expression that uses the number.
         *
         *  An example:
         *
         *  Initial control graph:
         *
         *     SeedPRNG -----> E_RandomNumber(Assign) ---> E_Add
         *         |
         *         |
         *         ----------> E_RandomNumber(Assign) ---> E_Mul
         *
         *
         *  After transformation:
         *
         *     SeedPRNG -----> Block  ---> E_Add
         *         |             |
         *         |           [Body]
         *         |             |
         *         |             v
         *         |           Assign (seedVGPR=RandomNumber(seedVGPR)) ---> E_RandomNumber
         *         |
         *         |
         *         ----------> Block  ---> E_Mul
         *                       |
         *                     [Body]
         *                       |
         *                       v
         *                     Assign (seedVGPR=RandomNumber(seedVGPR)) ---> E_RandomNumber
         */
        void updateSeed(KernelGraph& kgraph, std::unordered_map<int, std::vector<int>>& mapping)
        {
            auto& control = kgraph.control;
            auto& mapper  = kgraph.mapper;

            for(auto& [assignOp, prngs] : mapping)
            {
                auto blockOp = kgraph.control.addElement(Block());

                reconnect<Graph::Direction::Upstream>(kgraph, blockOp, assignOp);
                reconnect<Graph::Direction::Downstream>(kgraph, blockOp, assignOp);

                for(auto prng : prngs)
                {
                    auto seedVGPR
                        = mapper.get(prng, NaryArgument::DEST); // seedVGPR stores the seed value

                    auto dataFlowTag
                        = std::make_shared<Expression::Expression>(Expression::DataFlowTag{
                            seedVGPR, Register::Type::Vector, DataType::UInt32});
                    auto updateSeedExpr = std::make_shared<Expression::Expression>(
                        Expression::RandomNumber{dataFlowTag});
                    auto updateSeedOp
                        = control.addElement(Assign{Register::Type::Vector, updateSeedExpr});

                    Expression::setComment(updateSeedExpr, "Generate a new random number");

                    mapper.connect(updateSeedOp, seedVGPR, NaryArgument::DEST);

                    control.addElement(Body(), {blockOp}, {updateSeedOp});
                    control.addElement(Sequence(), {updateSeedOp}, {assignOp});
                }
            }
        }

        /** This transformation does:
         *
         *    1. Add TID subgraph to Coordinate graph is a random number needs to add TID to its seed.
         *
         *    2. Generate a new random number and store it in the seed VGPR of random number generator.
         *       This ensures all random numbers are different.
         */
        KernelGraph AddPRNG::apply(KernelGraph const& original)
        {
            auto kgraph = original;

            auto& coordinates = kgraph.coordinates;
            auto& control     = kgraph.control;
            auto& mapper      = kgraph.mapper;

            int tidVGPR = -1;

            // If a PRNG needs to add TID to the initial seed value, prepend a TID subgraph to
            // the PRNG's VGPR.
            for(auto seedPRNG : control.getNodes<SeedPRNG>())
            {
                bool const needTID = control.getNode<SeedPRNG>(seedPRNG).addTID;
                if(needTID)
                {
                    if(tidVGPR == -1)
                    {
                        // If tidVGPR (TID subgraph) not exists yet, create one and
                        // later this tidVGPR can be re-used by different random
                        // number generators that need to add TID
                        tidVGPR = createTIDSubgraph(m_context, coordinates);
                    }

                    auto seedVGPR = mapper.get(
                        seedPRNG, NaryArgument::DEST); // seedVGPR stores the seed value
                    coordinates.addElement(DataFlow(), {tidVGPR}, {seedVGPR});

                    // Map tidVGPR to the LHS arg of SeedPRNG op. RHS arg is the seed value given from user.
                    // Later on we will get VGPRs of LHS and RHS to set the seed of random number generator.
                    mapper.connect(seedPRNG, tidVGPR, NaryArgument::LHS);
                }
            }

            auto mapping = getMappingOfRandomNumbersAndPRNGs(kgraph);
            updateSeed(kgraph, mapping);

            return kgraph;
        }
    }
}