// Copyright (c) 2022, 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 <OgnMakeVector4Database.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <omni/graph/core/ogn/UsdTypes.h>
#include <fstream>
#include <iomanip>

namespace omni
{
namespace graph
{
namespace nodes
{

namespace
{
template <typename Type>
bool tryMakeVector(OgnMakeVector4Database& db)
{
    const auto x = db.inputs.x().template get<Type>();
    const auto y = db.inputs.y().template get<Type>();
    const auto z = db.inputs.z().template get<Type>();
    const auto w = db.inputs.w().template get<Type>();

    auto vector = db.outputs.tuple().template get<Type[4]>();

    if (vector && x && y && z && w)
    {
        (*vector)[0] = *x;
        (*vector)[1] = *y;
        (*vector)[2] = *z;
        (*vector)[3] = *w;
        return true;
    }

    const auto xArray = db.inputs.x().template get<Type[]>();
    const auto yArray = db.inputs.y().template get<Type[]>();
    const auto zArray = db.inputs.z().template get<Type[]>();
    const auto wArray = db.inputs.w().template get<Type[]>();

    auto vectorArray = db.outputs.tuple().template get<Type[][4]>();

    if (!vectorArray || !xArray || !yArray || !zArray || !wArray)
    {
        return false;
    }

    if (xArray->size() != yArray->size() || xArray->size() != zArray->size() || xArray->size() != wArray->size())
    {
        throw ogn::compute::InputError("Input arrays of different lengths x:" +
                        std::to_string(xArray->size()) + ", y:" +
                        std::to_string(yArray->size()) + ", z:" +
                        std::to_string(zArray->size()) + ", w:" +
                        std::to_string(wArray->size()));
    }

    vectorArray->resize(xArray->size());
    for (size_t i = 0; i < vectorArray->size(); i++)
    {
        (*vectorArray)[i][0] = (*xArray)[i];
        (*vectorArray)[i][1] = (*yArray)[i];
        (*vectorArray)[i][2] = (*zArray)[i];
        (*vectorArray)[i][3] = (*wArray)[i];
    }

    return true;
}
} // namespace
// Node to merge 4 scalers together to make 4-vector
class OgnMakeVector4
{
public:
    static bool compute(OgnMakeVector4Database& db)
    {
        // Compute the components, if the types are all resolved.
        try
        {
            if (tryMakeVector<double>(db))
                return true;
            else if (tryMakeVector<float>(db))
                return true;
            else if (tryMakeVector<pxr::GfHalf>(db))
                return true;
            else if (tryMakeVector<int32_t>(db))
                return true;
            else
            {
                db.logError("Failed to resolve input types");
                return false;
            }
        }
        catch (const std::exception& e)
        {
            db.logError("Vector could not be made: %s", e.what());
            return false;
        }

        return true;
    }

    static void onConnectionTypeResolve(const NodeObj& nodeObj)
    {
        auto x = nodeObj.iNode->getAttributeByToken(nodeObj, inputs::x.token());
        auto y = nodeObj.iNode->getAttributeByToken(nodeObj, inputs::y.token());
        auto z = nodeObj.iNode->getAttributeByToken(nodeObj, inputs::z.token());
        auto w = nodeObj.iNode->getAttributeByToken(nodeObj, inputs::w.token());

        auto vector = nodeObj.iNode->getAttributeByToken(nodeObj, outputs::tuple.token());
        auto xType = vector.iAttribute->getResolvedType(x);
        auto yType = vector.iAttribute->getResolvedType(y);
        auto zType = vector.iAttribute->getResolvedType(z);
        auto wType = vector.iAttribute->getResolvedType(w);

        std::array<AttributeObj, 4> attrs{ x, y, z, w };
        if (nodeObj.iNode->resolveCoupledAttributes(nodeObj, attrs.data(), attrs.size()))
        {
            xType = vector.iAttribute->getResolvedType(x);
            yType = vector.iAttribute->getResolvedType(y);
            zType = vector.iAttribute->getResolvedType(z);
            wType = vector.iAttribute->getResolvedType(w);
        }

        // Require inputs to be resolved before determining outputs' type
        if (xType.baseType != BaseDataType::eUnknown && yType.baseType != BaseDataType::eUnknown &&
            zType.baseType != BaseDataType::eUnknown && wType.baseType != BaseDataType::eUnknown)
        {
            std::array<AttributeObj, 5> attrs{ x, y, z, w, vector };
            std::array<uint8_t, 5> tuples{ 1, 1, 1, 1, 4 };
            std::array<uint8_t, 5> arrays{
                xType.arrayDepth,
                yType.arrayDepth,
                zType.arrayDepth,
                wType.arrayDepth,
                xType.arrayDepth
            };
            std::array<AttributeRole, 5> roles{ xType.role, yType.role, zType.role, wType.role, AttributeRole::eNone };
            nodeObj.iNode->resolvePartiallyCoupledAttributes(
                nodeObj, attrs.data(), tuples.data(), arrays.data(), roles.data(), attrs.size());
        }
    }
};

REGISTER_OGN_NODE();

}
}
}
