/*******************************************************************************
 *
 * MIT License
 *
 * Copyright (C) 2022-2026 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 *******************************************************************************/

#pragma once

#include <cassert>

#include <algorithm>
#include <cstddef>
#include <iomanip>
#include <iostream>
#include <numeric>
#include <vector>

#include <Tensile/DataTypes.hpp>
#include <Tensile/Debug.hpp>
#include <tensilelitehost/export.h>
#include <Tensile/Utils.hpp>

namespace TensileLite
{
    template <typename SizeIter>
    inline size_t CoordCount(SizeIter sizeBegin, SizeIter sizeEnd)
    {
        size_t rv = 1;

        while(sizeBegin != sizeEnd)
        {
            rv *= *sizeBegin;
            sizeBegin++;
        }

        return rv;
    }

    template <typename CoordIter, typename SizeIter>
    inline size_t CoordFlattenIndex(CoordIter coordBegin,
                                    CoordIter coordEnd,
                                    SizeIter  sizeBegin,
                                    SizeIter  sizeEnd)
    {
        auto   coord = coordEnd - 1;
        auto   size  = sizeEnd - 2;
        size_t rv    = 0;

        while(coord != coordBegin)
        {
            rv = (rv + *coord) * (*size);

            coord--;
            size--;
        }

        rv += *coord;

        return rv;
    }

    template <typename CoordIter, typename SizeIter>
    inline void CoordNumbered(
        size_t num, CoordIter coordBegin, CoordIter coordEnd, SizeIter sizeBegin, SizeIter sizeEnd)
    {
        auto coord = coordBegin;
        auto size  = sizeBegin;

        while(coord != coordEnd && size != sizeEnd)
        {
            *coord = num % *size;
            num /= *size;

            coord++;
            size++;
        }

        if(coord != coordEnd || size != sizeEnd)
            throw std::runtime_error("Inconsistent size of coordinates.");
    }

    template <typename CoordIter, typename SizeIter>
    inline void CoordNumberedExclude(size_t    num,
                                     CoordIter coordBegin,
                                     CoordIter coordEnd,
                                     SizeIter  sizeBegin,
                                     SizeIter  sizeEnd,
                                     size_t    excludeDim)
    {
        auto   coord      = coordBegin;
        auto   size       = sizeBegin;
        size_t currentDim = 0;
        while(coord != coordEnd && size != sizeEnd)
        {
            if(currentDim != excludeDim)
            {
                *coord = num % *size;
                num /= *size;
            }
            currentDim++;
            coord++;
            size++;
        }

        if(coord != coordEnd || size != sizeEnd)
            throw std::runtime_error("Inconsistent size of coordinates.");
    }

    template <typename CoordIter, typename SizeIter>
    inline bool IncrementCoord(CoordIter coordBegin,
                               CoordIter coordEnd,
                               SizeIter  sizeBegin,
                               SizeIter  sizeEnd)
    {
        auto coord = coordBegin;
        auto size  = sizeBegin;

        while(coord != coordEnd)
        {
            (*coord)++;
            if(*coord < *size)
                return true;

            *coord = 0;

            coord++;
            size++;
        }

        return false;
    }

    /**
 * Describes a tensor including dimensions, memory layout, and data type.
 * Decoupled from any particular pointer value or memory location.
 *
 * Provides functions for indexing and otherwise iterating through a tensor.
 */
    class TENSILELITEHOST_EXPORT TensorDescriptor
    {
    public:
        static const size_t UseDefaultStride;

        TensorDescriptor() {}

        TensorDescriptor(const char* name)
            : m_name(name)
        {
        }

        template <typename IterA, typename IterB>
        TensorDescriptor(const char*      name,
                         rocisa::DataType t,
                         IterA            sizesBegin,
                         IterA            sizesEnd,
                         IterB            stridesBegin,
                         IterB            stridesEnd)
            : m_name(name)
            , m_sizes(sizesBegin, sizesEnd)
            , m_strides(stridesBegin, stridesEnd)
            , m_dataType(t)
        {
            this->calculate();
        }

        template <typename Iter>
        TensorDescriptor(const char* name, rocisa::DataType t, Iter sizesBegin, Iter sizesEnd)
            : m_name(name)
            , m_sizes(sizesBegin, sizesEnd)
            , m_dataType(t)
        {
            this->calculate();
        }

        TensorDescriptor(const char* name, rocisa::DataType t, std::initializer_list<size_t> sizes)
            : m_name(name)
            , m_sizes(sizes)
            , m_dataType(t)

        {
            this->calculate();
        }

        TensorDescriptor(const char*                   name,
                         rocisa::DataType              t,
                         std::initializer_list<size_t> sizes,
                         std::initializer_list<size_t> strides)
            : m_name(name)
            , m_sizes(sizes)
            , m_strides(strides)
            , m_dataType(t)
        {
            this->calculate();
        }

        TensorDescriptor(const TensorDescriptor& rhs)
            : m_name(rhs.m_name)
            , m_sizes(rhs.m_sizes)
            , m_strides(rhs.m_strides)
            , m_totalLogicalElements(rhs.m_totalLogicalElements)
            , m_totalAllocatedElements(rhs.m_totalAllocatedElements)
            , m_dataType(rhs.m_dataType)
            , m_isOutput(rhs.m_isOutput)
        {
        }

        void resize(std::initializer_list<size_t>& sizes, std::initializer_list<size_t>& strides)
        {
            m_sizes   = sizes;
            m_strides = strides;
            this->calculate();
        }

        void setName(const char* name)
        {
            m_name = name;
        }

        const std::string& getName() const
        {
            return m_name;
        }

        void setAsOutput(bool isOutput)
        {
            m_isOutput = isOutput;
        }

        bool isOutput() const
        {
            return m_isOutput;
        }

        inline void calculate()
        {
            if(m_sizes.empty())
            {
                m_strides                = m_sizes;
                m_totalLogicalElements   = 0;
                m_totalAllocatedElements = 0;
                return;
            }

            m_strides.resize(m_sizes.size(), UseDefaultStride);
            if(m_strides[0] == UseDefaultStride)
            {
                m_strides[0] = 1;
            }
            m_totalLogicalElements = m_sizes[0];

            for(int i = 1; i < m_sizes.size(); i++)
            {
                m_totalLogicalElements *= m_sizes[i];

                if(m_strides[i] == UseDefaultStride)
                {
                    m_strides[i] = m_strides[i - 1] * m_sizes[i - 1];
                }
            }

            m_totalAllocatedElements = 1;
            for(int i = 0; i < m_sizes.size(); i++)
                m_totalAllocatedElements += m_strides[i] * (m_sizes[i] - 1);

            if(Debug::Instance().printTensorInfo())
            {
#ifndef CODE_COVERAGE
                std::cout << "TensorDescriptor:calculate  " << *this
                          << "totalLogicalElements=" << m_totalLogicalElements
                          << " totalAllocatedElem=" << m_totalAllocatedElements << "\n";
#endif
            }
        }

        const std::vector<size_t>& sizes() const
        {
            return m_sizes;
        }
        const std::vector<size_t>& strides() const
        {
            return m_strides;
        }
        bool empty() const
        {
            return m_sizes.empty();
        }

        void appendDim(size_t logicalCount);
        void appendDim(size_t logicalCount, size_t allocatedCount);

        /**
   * Returns the number of elements of padding in the given dimension (0 if
   * unpadded). May be negative if stride is less than size
   */
        int64_t dimensionPadding(size_t dim) const;

        /**
   * Collapses dimensions in the interval [begin, end).
   *
   * preconditions:
   * - end >= begin
   * - begin < dimensions()
   * - end <= dimensions()
   * - dimensions in the interval [begin, end-1) are not padded.
   *
   * postconditions:
   * - dimensions() is diminished by end-begin
   * - total elements (allocated and logical) remain the same
   * - dimension 'begin' is the product of all the dimensions in the interval
   * [begin, end).
   */
        void collapseDims(size_t begin, size_t end);

        size_t dimensions() const
        {
            return m_sizes.size();
        }
        size_t totalLogicalElements() const
        {
            return m_totalLogicalElements;
        }
        size_t totalAllocatedElements() const
        {
            return m_totalAllocatedElements;
        }
        size_t totalAllocatedBytes() const
        {
            // info.elementSize is the per-element segment size in bytes
            // (post-PR #6499: ElementSize = sizeof(T) / Packing). The
            // total byte count is integral as long as the element count
            // is a whole number of packed containers.
            auto const info = DataTypeInfo::Get(m_dataType);
            assert(totalAllocatedElements() % info.packing == 0);
            return multiplyElementSize(totalAllocatedElements(), info.elementSize);
        }

        float elementBytes() const
        {
            // info.elementSize is already the per-element segment size in
            // bytes (post-PR #6499: BaseTypeInfo::ElementSize = sizeof(T) /
            // Packing), so this returns 0.5 for Float4 and 0.75 for
            // Float6/BFloat6. Callers should pair this with
            // multiplyElementSize() from Utils.hpp when converting element
            // counts to byte counts.
            return DataTypeInfo::Get(m_dataType).elementSize;
        }

        void setDataType(rocisa::DataType type)
        {
            m_dataType = type;
        }

        rocisa::DataType dataType() const
        {
            return m_dataType;
        }

        template <typename Container>
        inline size_t index(Container const& indices) const
        {
            if(indices.size() != dimensions())
                throw std::runtime_error("Incorrect number of indices.");

            for(int i = 0; i < indices.size(); i++)
                if(indices[i] >= m_sizes[i])
                    throw std::runtime_error("Index out of bounds.");

            return std::inner_product(indices.begin(), indices.end(), m_strides.begin(), (size_t)0);
        }

        template <typename T>
        inline size_t index(std::initializer_list<T> indices) const
        {
            if(indices.size() != dimensions())
                throw std::runtime_error("Incorrect number of indices.");

            for(auto i = std::make_pair(indices.begin(), m_sizes.begin()); i.first != indices.end();
                i.first++, i.second++)
                if(*i.first >= *i.second)
                    throw std::runtime_error("Index out of bounds.");

            return std::inner_product(indices.begin(), indices.end(), m_strides.begin(), (size_t)0);
        }

        template <class... Ts,
                  typename = typename std::enable_if<
                      std::is_integral<typename std::common_type<Ts...>::type>::value>::type>
        inline size_t index(Ts... is) const
        {
            return this->index({is...});
        }

        inline bool incrementCoord(std::vector<size_t>& coord, size_t firstDimension = 0) const
        {
            if(coord.size() != dimensions())
                throw std::runtime_error(concatenate(
                    "Invalid coordinate size ", coord.size(), " for ", dimensions(), "-tensor"));

            if(firstDimension >= dimensions())
                return false;

            return IncrementCoord(
                coord.begin() + firstDimension, coord.end(), m_sizes.begin(), m_sizes.end());
        }

        bool operator==(const TensorDescriptor& rhs) const;
        bool operator!=(const TensorDescriptor& rhs) const;

        std::string ToString() const;

        friend TENSILELITEHOST_EXPORT std::ostream& operator<<(std::ostream&           stream,
                                                               const TensorDescriptor& t);

    private:
        std::string         m_name;
        std::vector<size_t> m_sizes;
        std::vector<size_t> m_strides;

        size_t m_totalLogicalElements   = 0;
        size_t m_totalAllocatedElements = 0;

        rocisa::DataType m_dataType = rocisa::DataType::Float;

        bool m_isOutput = false;
    };

    TENSILELITEHOST_EXPORT std::ostream& operator<<(std::ostream& stream, const TensorDescriptor& t);

    template <typename T>
    void WriteTensor1D(std::ostream&           stream,
                       T*                      data,
                       TensorDescriptor const& desc,
                       bool                    decorated = true)
    {
        if(desc.dimensions() != 1)
            throw std::runtime_error(
                "WriteTensor1D is only compatible with 1-dimensional tensors.");

        if(decorated)
            stream << "[";

        constexpr size_t packing = TypeInfo<std::remove_cv_t<T>>::Packing;

        if(desc.sizes()[0] > 0)
            stream << data[0];

        for(size_t i = packing; i < desc.sizes()[0]; i += packing)
            stream << " " << data[i / packing];

        if(decorated)
            stream << "]" << std::endl;
    }

    /**
 *  @brief Writes a tensor to an output stream.
 *
 * \param stream The stream to write to
 * \param data Pointer to the tensor data
 * \param desc Tensor descriptor
 * \param ptrValue Pointer value to print to describe the location of the data.
 * \param decorated Print brackets [] to indicate start/end of tensor dims
 */
    template <typename T>
    void WriteTensor(std::ostream&           stream,
                     T const*                data,
                     TensorDescriptor const& desc,
                     T const*                ptrValue  = nullptr,
                     bool                    decorated = true)
    {
        stream << "Tensor(";
        streamJoin(stream, desc.sizes(), ", ");
        stream << ", data_ptr: " << data << ")" << std::endl;

        if(desc.dimensions() == 0)
            return;

        if(desc.dimensions() == 1)
        {
            WriteTensor1D(stream, data, desc, decorated);
            return;
        }

        auto const&         sizes = desc.sizes();
        std::vector<size_t> coord(desc.dimensions(), 0);
        const auto          stride0 = desc.strides()[0];

        auto upperDimCount = CoordCount(sizes.begin() + 2, sizes.end());

        for(size_t idx = 0; idx < upperDimCount; idx++)
        {
            CoordNumbered(idx, coord.begin() + 2, coord.end(), sizes.begin() + 2, sizes.end());

            coord[0] = 0;
            coord[1] = 0;

            if(decorated)
            {
                stream << "(";
                streamJoin(stream, coord, ", ");
                stream << ")" << std::endl << "[" << std::endl;
            }

            for(coord[1] = 0; coord[1] < sizes[1]; coord[1]++)
            {
                coord[0] = 0;

                auto const* localPtr = data + (desc.index(coord) / TypeInfo<T>::Packing);

                if(sizes[0] > 0)
                    stream << localPtr[0];

                for(coord[0] = TypeInfo<T>::Packing; coord[0] < sizes[0]; coord[0]+=TypeInfo<T>::Packing)
                {
                    stream << " " << localPtr[coord[0] * stride0 / TypeInfo<T>::Packing];
                }

                stream << std::endl;
            }

            if(decorated)
            {
                stream << "]" << std::endl;
            }
        }
    }

} // namespace TensileLite