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

#pragma once

#include <cstdint>
#include <map>
#include <memory>
#include <string>
#include <vector>

#include <rocRoller/Expression.hpp>
#include <rocRoller/Graph/Hypergraph.hpp>
#include <rocRoller/KernelGraph/CoordinateGraph/CoordinateEdge.hpp>
#include <rocRoller/KernelGraph/CoordinateGraph/Dimension.hpp>

namespace rocRoller
{
    /**
     * Coordinate transform (index calculations) routines.
     *
     * Coordinate transforms are represented as graphs.  A coordinate
     * transform graph encodes:
     *
     * - data flow
     * - data locality
     * - how indexes are computed/transformed
     *
     * Given the geometry of a tensor (sizes, strides), the graph can
     * be used to determine index values of given working coordinates.
     *
     * The graph answers the question:
     * What is the relationship between the dimensions of:
     *  - Input and output tensors
     *  - GPU execution constructs such as workgroups, workitems, wavefronts
     *  - Software constructs such as loops & unrolls
     *  - Temporary storage such as registers (V/SGPRs), LDS, scratch space
     *
     * It maps problem memory space to the hardware aspects, like
     * workgroups, workitems, VGPRs, LDS, and software constraints,
     * like loop dimensions/unrolls.
     *
     * Nodes in the graph represent abstract "dimensions".  These can
     * represent, for example: tensors provided by the user, linear
     * arrays stored in LDS, or loop indexes.
     *
     * Edges in the graph represent how coordinates are transformed.
     *
     * Throughout this documentation we use the following notation to
     * describe coordinate transforms: consider a "flatten" transform
     * with input dimensions
     *
     *   I = Dimension(size=n_i, stride=s_i)
     *   J = Dimension(size=n_j, stride=s_j)
     *
     * and output dimension
     *
     *   O = Dimension().
     *
     * The forward coordinate transform is denoted
     *
     *   Flatten(I, J; O)(i, j) = i * n_j + j
     *
     * and the inverse coordinate transform is denoted
     *
     *   Flatten'(O; I, J)(o) = { o / n_j, o % n_j }.
     *
     * That is,
     *
     *   Flatten(input dimensions; output dimensions)(input indexes)
     *
     * and inverse
     *
     *   Flatten'(output dimensions; input dimensions)(output indexes)
     *
     *
     * Nodes may have negative tag values, to indicate unallocated tags.
     * When the node is added to a graph (control or coordinate), it will
     * be replaced with a positive value that is otherwise unused in that
     * graph.
     *
     */
    namespace KernelGraph::CoordinateGraph
    {
        /**
         * Coordinate-transform HyperGraph.
         *
         * Nodes in the graph represent single dimensions (or
         * coordinates).
         *
         * Hyper-edges describe how to transform coordinates and/or
         * apply operations.
         */
        class CoordinateGraph : public Graph::Hypergraph<Dimension, Edge>
        {
            bool m_changesRestricted = false;

        public:
            using Base = Graph::Hypergraph<Dimension, Edge>;

            CoordinateGraph() = default;

            std::vector<Expression::ExpressionPtr>
                forward(std::vector<Expression::ExpressionPtr> sdims,
                        std::vector<int> const&                srcs,
                        std::vector<int> const&                dsts) const;

            std::vector<Expression::ExpressionPtr>
                reverse(std::vector<Expression::ExpressionPtr> sdims,
                        std::vector<int> const&                srcs,
                        std::vector<int> const&                dsts) const;

            EdgeType getEdgeType(int index) const;

            template <Graph::Direction Dir, typename Visitor>
            std::vector<Expression::ExpressionPtr>
                traverse(std::vector<Expression::ExpressionPtr> sdims,
                         std::vector<int> const&                srcs,
                         std::vector<int> const&                dsts,
                         Visitor&                               visitor) const;

            template <Graph::Direction Dir>
            bool hasPath(std::vector<int> const& srcs, std::vector<int> const& dsts) const;

            /**
             * @brief Get a node/edge from the coordinate graph.
             *
             * If the element specified by tag cannot be converted to
             * T, the return value is empty.
             *
             * @param tag Graph tag/index.
             */
            template <typename T>
            requires(std::constructible_from<CoordinateGraph::Element, T>) std::optional<T> get(
                int tag)
            const;

            /**
             *  Check if modifying an element (index) is allowed or not. This
             *  only comes into effect when the graph is in restricted mode.
             */
            virtual bool isModificationAllowed(int index) const override;

            /**
             *  Set the graph to be in restricted mode. Some operations would
             *  be disallowed when in restricted mode.
             */
            void setRestricted()
            {
                m_changesRestricted = true;
            }
        };

        std::string name(CoordinateGraph::Element const& el);
    }
}

#include <rocRoller/KernelGraph/CoordinateGraph/CoordinateGraph_impl.hpp>
