// 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 <OgnMakeVector3Database.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(OgnMakeVector3Database& db, size_t count = 1)
{
    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 vector = db.outputs.tuple().template get<Type[3]>();

    if (x && y && z && vector)
    {
        const auto px = x.vectorized(count);
        const auto py = y.vectorized(count);
        const auto pz = z.vectorized(count);
        const auto pvector = vector.vectorized(count);

        if (not(pvector.empty() || px.empty() || py.empty() || pz.empty()))
        {
            for (size_t i = 0; i < count; i++)
            {
                pvector[i][0] = px[i];
                pvector[i][1] = py[i];
                pvector[i][2] = pz[i];
            }
            return true;
        }
        return false;
    }

    for(size_t i = 0; i < count ; ++i)
    {
        const auto xArray = db.inputs.x(i).template get<Type[]>();
        const auto yArray = db.inputs.y(i).template get<Type[]>();
        const auto zArray = db.inputs.z(i).template get<Type[]>();

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

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

    
        if (xArray->size() != yArray->size() || xArray->size() != zArray->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()));
        }

        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];
        }
    }
 
    return true;
}
} // namespace
// Node to merge 3 scalers together to make 3-vector
class OgnMakeVector3
{
public:
    
    static size_t computeVectorized(OgnMakeVector3Database& db, size_t count)
    {
        // Compute the components, if the types are all resolved.
        try
        {
            if (tryMakeVector<double>(db, count))
                return count;
            else if (tryMakeVector<float>(db, count))
                return count;
            else if (tryMakeVector<pxr::GfHalf>(db, count))
                return count;
            else if (tryMakeVector<int32_t>(db, count))
                return count;
            else
            {
                db.logError("Failed to resolve input types");
                return 0;
            }
        }
        catch (const std::exception& e)
        {
            db.logError("Vector could not be made: %s", e.what());
            return 0;
        }

        return 0;
    }

    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 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);

        // If one of the inputs is resolved we can resolve the other because they should match
        std::array<AttributeObj, 3> attrs { x, y, z };
        if (nodeObj.iNode->resolveCoupledAttributes(nodeObj, attrs.data(), attrs.size()))
        {
            xType = vector.iAttribute->getResolvedType(x);
            yType = vector.iAttribute->getResolvedType(y);
            zType = vector.iAttribute->getResolvedType(z);
        }

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

REGISTER_OGN_NODE();

}
}
}
