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

#include <rocRoller/KernelGraph/KernelGraph.hpp>
#include <rocRoller/KernelGraph/Transforms/AddConvert.hpp>
#include <rocRoller/KernelGraph/Utils.hpp>

#include <rocRoller/Utilities/Concepts.hpp>

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

        class AddConvertOperations
        {
        public:
            AddConvertOperations(ContextPtr context)
                : m_context(context)
            {
            }

            void addConverts(KernelGraph& graph);

        private:
            struct ConvertLocation
            {
                /// Storage coordinate (segmented) that needs to be converted
                int storageCoord;

                /// Load operations that write to storageCoord
                std::unordered_set<int> loadOps;

                /// Multiply operations that read from storageCoord
                std::vector<std::pair<int, NaryArgument>> uses;
            };

            void stageMultiplyConverts(KernelGraph const& graph);
            void commit(KernelGraph& graph);

            // Map from storage coordinate to vector of {Multiply operation tag, LHS/RHS} pairs.
            std::map<int, std::vector<std::pair<int, NaryArgument>>> m_multiplyArgs;
            // Map from storage coordinate to set of LoadTiled/LoadLDSTile operations
            std::map<int, std::unordered_set<int>> m_loadMap;
            // Map from storage coordinate to DataType
            std::map<int, DataType> m_storageDataType;

            std::vector<ConvertLocation> m_locations;

            ContextPtr m_context;
        };

        void AddConvertOperations::stageMultiplyConverts(KernelGraph const& graph)
        {
            // Populate data structures with information about multiplies
            // and loads
            auto root = graph.control.roots().only();
            AssertFatal(root.has_value());

            auto allNodes = graph.control.depthFirstVisit(*root).filter(
                graph.control.isElemType<Operation>());

            for(const auto node : allNodes)
            {
                auto visitor = rocRoller::overloaded{
                    [&](Multiply op) {
                        auto [macATag, macA] = graph.getDimension<MacroTile>(
                            node, Connections::typeArgument<MacroTile>(NaryArgument::LHS));
                        m_multiplyArgs[macATag].push_back({node, NaryArgument::LHS});

                        auto [macBTag, macB] = graph.getDimension<MacroTile>(
                            node, Connections::typeArgument<MacroTile>(NaryArgument::RHS));
                        m_multiplyArgs[macBTag].push_back({node, NaryArgument::RHS});

                        if(m_context->targetArchitecture().HasCapability(
                               GPUCapability::PartiallyActiveWaveSize))
                        {
                            if(op.scaleA == Operations::ScaleMode::Separate)
                            {
                                auto [macScaleATag, macScaleA] = graph.getDimension<MacroTile>(
                                    node,
                                    Connections::typeArgument<MacroTile>(NaryArgument::LHS_SCALE));
                                m_multiplyArgs[macScaleATag].push_back(
                                    {node, NaryArgument::LHS_SCALE});
                            }

                            if(op.scaleB == Operations::ScaleMode::Separate)
                            {
                                auto [macScaleBTag, macScaleB] = graph.getDimension<MacroTile>(
                                    node,
                                    Connections::typeArgument<MacroTile>(NaryArgument::RHS_SCALE));
                                m_multiplyArgs[macScaleBTag].push_back(
                                    {node, NaryArgument::RHS_SCALE});
                            }
                        }
                    },
                    [&](CIsAnyOf<LoadTiled, LoadLDSTile> auto op) {
                        auto coord = graph.mapper.get<MacroTile>(node);
                        m_loadMap[coord].insert(node);
                        m_storageDataType[graph.mapper.get<MacroTile>(node)]
                            = getDataType(graph.control.getNode(node));
                        Log::debug("(node {}) m_storageDataType[{}] = {}",
                                   node,
                                   graph.mapper.get<MacroTile>(node),
                                   toString(getDataType(graph.control.getNode(node))));
                    },
                    [&](auto op) {}};

                std::visit(visitor, graph.control.getNode(node));
            }

            for(auto const& [storage, multiplies] : m_multiplyArgs)
            {
                auto packed = DataTypeInfo::Get(m_storageDataType[storage]).packedVariableType();
                if(!packed)
                    continue;
                if(m_storageDataType[storage] == packed->dataType)
                    continue;

                Log::debug("{}: convert {} to {}",
                           storage,
                           toString(m_storageDataType[storage]),
                           toString(*packed));

                m_locations.emplace_back(storage, m_loadMap[storage], multiplies);
            }
        }

        void AddConvertOperations::commit(KernelGraph& graph)
        {
            namespace CT = rocRoller::KernelGraph::CoordinateGraph;

            std::map<int, int> newStorageDuplicate;

            for(auto& location : m_locations)
            {
                // Create new node for convert in control graph and storage in coordinate graph
                auto newStorage = graph.coordinates.addElement(
                    graph.coordinates.getElement(location.storageCoord));
                auto dataFlow = std::make_shared<Expression::Expression>(Expression::DataFlowTag{
                    location.storageCoord, Register::Type::Vector, DataType::None});
                auto packed   = DataTypeInfo::Get(m_storageDataType[location.storageCoord])
                                  .packedVariableType()
                                  ->dataType;

                for(auto loadOp : location.loadOps)
                {
                    auto convertNode = graph.control.addElement(
                        Assign{Register::Type::Vector, Expression::convert(packed, dataFlow), 0});
                    graph.mapper.connect(convertNode, newStorage, NaryArgument::DEST);
                    insertAfter(graph, loadOp, convertNode, convertNode);

                    Log::debug("{}: added convert {} (new storage {})",
                               location.storageCoord,
                               convertNode,
                               newStorage);
                }

                // Change mappings for all multiplies
                for(auto const& [node, arg] : location.uses)
                {
                    graph.mapper.disconnect(
                        node, location.storageCoord, Connections::typeArgument<MacroTile>(arg));
                    graph.mapper.connect(
                        node, newStorage, Connections::typeArgument<MacroTile>(arg));
                }

                // Add newStorage to coordinate graph, keeping track of duplicates
                // If duplicate, find original.
                auto duplicateStorage = only(graph.coordinates.getOutputNodeIndices(
                    location.storageCoord, CT::isEdge<Duplicate>));

                auto originalStorage = duplicateStorage ? *duplicateStorage : location.storageCoord;

                // If original hasn't been seen, insert newStorage after original with DataFlow edge
                if(newStorageDuplicate.count(originalStorage) == 0)
                {
                    auto childEdges = graph.coordinates.getNeighbours<Graph::Direction::Downstream>(
                        originalStorage);
                    auto loc = graph.coordinates.getLocation(originalStorage);

                    for(auto const& edge : loc.outgoing)
                    {
                        if(!CT::isEdge<DataFlow>(graph.coordinates.getEdge(edge)))
                            continue;

                        auto             edgeLoc = graph.coordinates.getLocation(edge);
                        std::vector<int> newInputs;
                        for(auto const& input : edgeLoc.incoming)
                        {
                            if(input != originalStorage)
                                newInputs.push_back(input);
                            else
                                newInputs.push_back(newStorage);
                        }

                        graph.coordinates.addElement(DataFlow(), newInputs, edgeLoc.outgoing);
                        graph.coordinates.deleteElement(edge);
                    }
                    graph.coordinates.addElement(DataFlow(), {originalStorage}, {newStorage});
                    newStorageDuplicate[originalStorage] = newStorage;
                }
                // if original has been seen, add Duplicate edge from newStorage to saved node
                else
                {
                    graph.coordinates.addElement(
                        Duplicate(), {newStorage}, {newStorageDuplicate[originalStorage]});
                }
            }
        }

        void AddConvertOperations::addConverts(KernelGraph& graph)
        {
            stageMultiplyConverts(graph);
            commit(graph);
        }

        KernelGraph AddConvert::apply(KernelGraph const& k)
        {
            auto graph = k;

            AddConvertOperations adder{m_context};

            adder.addConverts(graph);

            return graph;
        }
    }
}