// Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto.  Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//
#include <OgnArrayRotateDatabase.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <omni/graph/core/ogn/string.h>
#include <omni/graph/core/ogn/Types.h>

#include <algorithm>
namespace omni
{
namespace graph
{
namespace nodes
{

namespace
{
// Custom implementation of std::rotate.
// Using std::rotate fails to build on Linux when the underlying type is a tuple (eg int[3])
template<class ForwardIt>
constexpr
ForwardIt
rotateArray(ForwardIt first, ForwardIt n_first, ForwardIt last)
{
   if(first == n_first) return last;
   if(n_first == last) return first;

   ForwardIt read      = n_first;
   ForwardIt write     = first;
   ForwardIt next_read = first; // read position for when "read" hits "last"

   while(read != last) {
      if(write == next_read) next_read = read; // track where "first" went
      std::iter_swap(write++, read++);
   }

   // rotate the remaining sequence into place
   (rotateArray)(write, next_read, last);
   return write;
}

template<typename T>
bool tryComputeAssumingType(OgnArrayRotateDatabase& db)
{
    const auto& steps = db.inputs.steps();

    if (auto array = db.inputs.array().template get<T[]>())
    {
        if (auto result = db.outputs.array().template get<T[]>())
        {
            *result = *array;
            if (!result->empty())
            {
                if (steps < 0)
                {
                    rotateArray(result->begin(), result->begin() - steps, result->end());
                }
                else if (steps > 0)
                {
                    rotateArray(result->rbegin(), result->rbegin() + steps, result->rend());
                }
            }
            return true;
        }
    }
    return false;
}

} // namespace

class OgnArrayRotate
{
public:
    static bool compute(OgnArrayRotateDatabase& db)
    {
        auto& inputType = db.inputs.array().type();
        try
        {
            switch (inputType.baseType)
            {
            case BaseDataType::eBool:
                return tryComputeAssumingType<bool>(db);
            case BaseDataType::eToken:
                return tryComputeAssumingType<ogn::Token>(db);
            case BaseDataType::eDouble:
                switch (inputType.componentCount)
                {
                    case 1:  return tryComputeAssumingType<double>(db);
                    case 2:  return tryComputeAssumingType<double[2]>(db);
                    case 3:  return tryComputeAssumingType<double[3]>(db);
                    case 4:  return tryComputeAssumingType<double[4]>(db);
                    case 9:  return tryComputeAssumingType<double[9]>(db);
                    case 16: return tryComputeAssumingType<double[16]>(db);
                    default: break;
                }
            case BaseDataType::eFloat:
                switch (inputType.componentCount)
                {
                    case 1: return tryComputeAssumingType<float>(db);
                    case 2: return tryComputeAssumingType<float[2]>(db);
                    case 3: return tryComputeAssumingType<float[3]>(db);
                    case 4: return tryComputeAssumingType<float[4]>(db);
                    default: break;
                }
            case BaseDataType::eHalf:
                switch (inputType.componentCount)
                {
                    case 1: return tryComputeAssumingType<pxr::GfHalf>(db);
                    case 2: return tryComputeAssumingType<pxr::GfHalf[2]>(db);
                    case 3: return tryComputeAssumingType<pxr::GfHalf[3]>(db);
                    case 4: return tryComputeAssumingType<pxr::GfHalf[4]>(db);
                    default: break;
                }
            case BaseDataType::eInt:
                switch (inputType.componentCount)
                {
                    case 1: return tryComputeAssumingType<int32_t>(db);
                    case 2: return tryComputeAssumingType<int32_t[2]>(db);
                    case 3: return tryComputeAssumingType<int32_t[3]>(db);
                    case 4: return tryComputeAssumingType<int32_t[4]>(db);
                    default: break;
                }
                ;
            case BaseDataType::eInt64:
                return tryComputeAssumingType<int64_t>(db);
            case BaseDataType::eUChar:
                return tryComputeAssumingType<unsigned char>(db);
            case BaseDataType::eUInt:
                return tryComputeAssumingType<uint32_t>(db);
            case BaseDataType::eUInt64:
                return tryComputeAssumingType<uint64_t>(db);
            default: break;
            }

            db.logWarning("Failed to resolve input types");
        }
        catch (ogn::compute::InputError &error)
        {
            db.logWarning("OgnArrayRotate: %s", error.what());
        }
        return false;
    }

    static void onConnectionTypeResolve(const NodeObj& node){
        std::array<AttributeObj, 2> attrs {
            node.iNode->getAttributeByToken(node, inputs::array.token()),
            node.iNode->getAttributeByToken(node, outputs::array.token())
        };
        node.iNode->resolveCoupledAttributes(node, attrs.data(), attrs.size());
    }
};

REGISTER_OGN_NODE()

} // namespace nodes
} // namespace graph
} // namespace omni
