/******************************************************************************

 * Name:     gdalmultiDim_meshgrid.cpp

 * Project:  GDAL Core

 * Purpose:  Return a vector of coordinate matrices from coordinate vectors.

 * Author:   Even Rouault <even.rouault at spatialys.com>

 *

 ******************************************************************************

 * Copyright (c) 2024, Even Rouault <even.rouault at spatialys.com>

 *

 * SPDX-License-Identifier: MIT

 ****************************************************************************/

#include "gdal_priv.h"

#include <algorithm>

#include <limits>

/************************************************************************/

/*                       GetConcatenatedNames()                         */

/************************************************************************/

static std::string

GetConcatenatedNames(const std::vector<std::shared_ptr<GDALMDArray>> &apoArrays)

{

    std::string ret;

    for (const auto &poArray : apoArrays)

    {

        if (!ret.empty())

            ret += ", ";

        ret += poArray->GetFullName();

    }

    return ret;

}

/************************************************************************/

/*                         GDALMDArrayMeshGrid                          */

/************************************************************************/

class GDALMDArrayMeshGrid final : public GDALMDArray

{

    const std::vector<std::shared_ptr<GDALMDArray>> m_apoArrays;

    std::vector<std::shared_ptr<GDALDimension>> m_apoDims{};

    const size_t m_iDim;

    const bool m_bIJIndexing;

  protected:

    explicit GDALMDArrayMeshGrid(

        const std::vector<std::shared_ptr<GDALMDArray>> &apoArrays,

        const std::vector<std::shared_ptr<GDALDimension>> &apoDims, size_t iDim,

        bool bIJIndexing)

        : GDALAbstractMDArray(std::string(),

                              "Mesh grid view of " +

                                  GetConcatenatedNames(apoArrays)),

          GDALMDArray(std::string(),

                      "Mesh grid view of " + GetConcatenatedNames(apoArrays)),

          m_apoArrays(apoArrays), m_apoDims(apoDims), m_iDim(iDim),

          m_bIJIndexing(bIJIndexing)

    {

    }

    bool IRead(const GUInt64 *arrayStartIdx, const size_t *count,

               const GInt64 *arrayStep, const GPtrDiff_t *bufferStride,

               const GDALExtendedDataType &bufferDataType,

               void *pDstBuffer) const override;

  public:

    static std::shared_ptr<GDALMDArrayMeshGrid>

    Create(const std::vector<std::shared_ptr<GDALMDArray>> &apoArrays,

           size_t iDim, bool bIJIndexing)

    {

        std::vector<std::shared_ptr<GDALDimension>> apoDims;

        for (size_t i = 0; i < apoArrays.size(); ++i)

        {

            const size_t iTranslatedDim = (!bIJIndexing && i <= 1) ? 1 - i : i;

            apoDims.push_back(apoArrays[iTranslatedDim]->GetDimensions()[0]);

        }

        auto newAr(std::shared_ptr<GDALMDArrayMeshGrid>(

            new GDALMDArrayMeshGrid(apoArrays, apoDims, iDim, bIJIndexing)));

        newAr->SetSelf(newAr);

        return newAr;

    }

    bool IsWritable() const override

    {

        return false;

    }

    const std::string &GetFilename() const override

    {

        return m_apoArrays[m_iDim]->GetFilename();

    }

    const std::vector<std::shared_ptr<GDALDimension>> &

    GetDimensions() const override

    {

        return m_apoDims;

    }

    const GDALExtendedDataType &GetDataType() const override

    {

        return m_apoArrays[m_iDim]->GetDataType();

    }

    std::shared_ptr<GDALAttribute>

    GetAttribute(const std::string &osName) const override

    {

        return m_apoArrays[m_iDim]->GetAttribute(osName);

    }

    std::vector<std::shared_ptr<GDALAttribute>>

    GetAttributes(CSLConstList papszOptions = nullptr) const override

    {

        return m_apoArrays[m_iDim]->GetAttributes(papszOptions);

    }

    const std::string &GetUnit() const override

    {

        return m_apoArrays[m_iDim]->GetUnit();

    }

    const void *GetRawNoDataValue() const override

    {

        return m_apoArrays[m_iDim]->GetRawNoDataValue();

    }

    double GetOffset(bool *pbHasOffset,

                     GDALDataType *peStorageType) const override

    {

        return m_apoArrays[m_iDim]->GetOffset(pbHasOffset, peStorageType);

    }

    double GetScale(bool *pbHasScale,

                    GDALDataType *peStorageType) const override

    {

        return m_apoArrays[m_iDim]->GetScale(pbHasScale, peStorageType);

    }

};

/************************************************************************/

/*                             IRead()                                  */

/************************************************************************/

bool GDALMDArrayMeshGrid::IRead(const GUInt64 *arrayStartIdx,

                                const size_t *count, const GInt64 *arrayStep,

                                const GPtrDiff_t *bufferStride,

                                const GDALExtendedDataType &bufferDataType,

                                void *pDstBuffer) const

{

    const size_t nBufferDTSize = bufferDataType.GetSize();

    const size_t iTranslatedDim =

        (!m_bIJIndexing && m_iDim <= 1) ? 1 - m_iDim : m_iDim;

    std::vector<GByte> abyTmpData(nBufferDTSize * count[iTranslatedDim]);

    const GPtrDiff_t strideOne[] = {1};

    if (!m_apoArrays[m_iDim]->Read(&arrayStartIdx[iTranslatedDim],

                                   &count[iTranslatedDim],

                                   &arrayStep[iTranslatedDim], strideOne,

                                   bufferDataType, abyTmpData.data()))

        return false;

    const auto nDims = GetDimensionCount();

    struct Stack

    {

        size_t nIters = 0;

        GByte *dst_ptr = nullptr;

        GPtrDiff_t dst_inc_offset = 0;

    };

    // +1 to avoid -Werror=null-dereference

    std::vector<Stack> stack(nDims + 1);

    for (size_t i = 0; i < nDims; i++)

    {

        stack[i].dst_inc_offset =

            static_cast<GPtrDiff_t>(bufferStride[i] * nBufferDTSize);

    }

    stack[0].dst_ptr = static_cast<GByte *>(pDstBuffer);

    size_t dimIdx = 0;

    size_t valIdx = 0;

lbl_next_depth:

    if (dimIdx == nDims - 1)

    {

        auto nIters = count[dimIdx];

        GByte *dst_ptr = stack[dimIdx].dst_ptr;

        if (dimIdx == iTranslatedDim)

        {

            valIdx = 0;

            while (true)

            {

                GDALExtendedDataType::CopyValue(

                    &abyTmpData[nBufferDTSize * valIdx], bufferDataType,

                    dst_ptr, bufferDataType);

                if ((--nIters) == 0)

                    break;

                ++valIdx;

                dst_ptr += stack[dimIdx].dst_inc_offset;

            }

        }

        else

        {

            while (true)

            {

                GDALExtendedDataType::CopyValue(

                    &abyTmpData[nBufferDTSize * valIdx], bufferDataType,

                    dst_ptr, bufferDataType);

                if ((--nIters) == 0)

                    break;

                dst_ptr += stack[dimIdx].dst_inc_offset;

            }

        }

    }

    else

    {

        if (dimIdx == iTranslatedDim)

            valIdx = 0;

        stack[dimIdx].nIters = count[dimIdx];

        while (true)

        {

            dimIdx++;

            stack[dimIdx].dst_ptr = stack[dimIdx - 1].dst_ptr;

            goto lbl_next_depth;

        lbl_return_to_caller:

            dimIdx--;

            if ((--stack[dimIdx].nIters) == 0)

                break;

            if (dimIdx == iTranslatedDim)

                ++valIdx;

            stack[dimIdx].dst_ptr += stack[dimIdx].dst_inc_offset;

        }

    }

    if (dimIdx > 0)

    {

        goto lbl_return_to_caller;

    }

    if (bufferDataType.NeedsFreeDynamicMemory())

    {

        for (size_t i = 0; i < count[iTranslatedDim]; ++i)

        {

            bufferDataType.FreeDynamicMemory(&abyTmpData[i * nBufferDTSize]);

        }

    }

    return true;

}

/************************************************************************/

/*                      GDALMDArrayGetMeshGrid()                        */

/************************************************************************/

/** Return a list of multidimensional arrays from a list of one-dimensional

 * arrays.

 *

 * This is typically used to transform one-dimensional longitude, latitude

 * arrays into 2D ones.

 *

 * More formally, for one-dimensional arrays x1, x2,..., xn with lengths

 * Ni=len(xi), returns (N1, N2, ..., Nn) shaped arrays if indexing="ij" or

 * (N2, N1, ..., Nn) shaped arrays if indexing="xy" with the elements of xi

 * repeated to fill the matrix along the first dimension for x1, the second

 * for x2 and so on.

 *

 * For example, if x = [1, 2], and y = [3, 4, 5],

 * GetMeshGrid([x, y], ["INDEXING=xy"]) will return [xm, ym] such that

 * xm=[[1, 2],[1, 2],[1, 2]] and ym=[[3, 3],[4, 4],[5, 5]],

 * or more generally xm[any index][i] = x[i] and ym[i][any index]=y[i]

 *

 * and

 * GetMeshGrid([x, y], ["INDEXING=ij"]) will return [xm, ym] such that

 * xm=[[1, 1, 1],[2, 2, 2]] and ym=[[3, 4, 5],[3, 4, 5]],

 * or more generally xm[i][any index] = x[i] and ym[any index][i]=y[i]

 *

 * The currently supported options are:

 * <ul>

 * <li>INDEXING=xy/ij: Cartesian ("xy", default) or matrix ("ij") indexing of

 * output.

 * </li>

 * </ul>

 *

 * This is the same as

 * <a href="https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html">numpy.meshgrid()</a>

 * function.

 *

 * This is the same as the C function GDALMDArrayGetMeshGrid()

 *

 * @param apoArrays Input arrays

 * @param papszOptions NULL, or NULL terminated list of options.

 *

 * @return an array of coordinate matrices

 * @since 3.10

 */

/* static */ std::vector<std::shared_ptr<GDALMDArray>> GDALMDArray::GetMeshGrid(

    const std::vector<std::shared_ptr<GDALMDArray>> &apoArrays,

    CSLConstList papszOptions)

{

    std::vector<std::shared_ptr<GDALMDArray>> ret;

    for (const auto &poArray : apoArrays)

    {

        if (poArray->GetDimensionCount() != 1)

        {

            CPLError(CE_Failure, CPLE_NotSupported,

                     "Only 1-D input arrays are accepted");

            return ret;

        }

    }

    const char *pszIndexing =

        CSLFetchNameValueDef(papszOptions, "INDEXING", "xy");

    if (!EQUAL(pszIndexing, "xy") && !EQUAL(pszIndexing, "ij"))

    {

        CPLError(CE_Failure, CPLE_NotSupported,

                 "Only INDEXING=xy or ij is accepted");

        return ret;

    }

    const bool bIJIndexing = EQUAL(pszIndexing, "ij");

    for (size_t i = 0; i < apoArrays.size(); ++i)

    {

        ret.push_back(GDALMDArrayMeshGrid::Create(apoArrays, i, bIJIndexing));

    }

    return ret;

}