// Copyright (c) 2021-2021, 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 <OgnGetMatrix4QuaternionDatabase.h>
#include <omni/math/linalg/matrix.h>
#include <omni/math/linalg/quat.h>
#include "XformUtils.h"

using omni::math::linalg::matrix4d;
using omni::math::linalg::quatd;

namespace omni
{
namespace graph
{
namespace nodes
{

namespace
{
    void getQuaternion(const double input[16], double result[4])
    {
        quatd quat = extractRotationQuatd(*reinterpret_cast<const matrix4d*>(input));
        // In fabric, quaternions are stored as (XYZW). Note the quatd constructor uses (WXYZ)s
        memcpy(&result[0], quat.GetImaginary().data(), sizeof(double) * 3);
        result[3] = quat.GetReal();
    }
}

class OgnGetMatrix4Quaternion
{
public:
    static bool compute(OgnGetMatrix4QuaternionDatabase& db)
    {
        const auto& matrixInput = db.inputs.matrix();
        auto quaternionOutput = db.outputs.quaternion();
        bool failed = true;

        // Singular case
        if (auto matrix = matrixInput.get<double[16]>())
        {           
            if (auto output = quaternionOutput.get<double[4]>())
            {
                getQuaternion(*matrix, *output);
                failed = false;
            }
        }
        // Array case
        else if (auto matrices = matrixInput.get<double[][16]>())
        {
            if (auto output = quaternionOutput.get<double[][4]>())
            {
                output->resize(matrices->size());
                for (size_t i = 0; i < matrices.size(); i++)
                {
                    getQuaternion((*matrices)[i], (*output)[i]);
                }
                failed = false;
            }
        }
        else
        {
            db.logError("Input type for matrix input not supported");
            return false;
        }
        
        if (failed)
        {
            db.logError("Input and output depths need to align: Matrix input with depth %s, output with depth %s", 
                    matrixInput.type().arrayDepth, quaternionOutput.type().arrayDepth);
            return false;
        }

        return true;
    }
    
    static void onConnectionTypeResolve(const NodeObj& node) {
        auto start = node.iNode->getAttribute(node, OgnGetMatrix4QuaternionAttributes::inputs::matrix.m_name);
        auto result = node.iNode->getAttribute(node, OgnGetMatrix4QuaternionAttributes::outputs::quaternion.m_name);
        auto startType = start.iAttribute->getResolvedType(start);

        std::array<AttributeObj, 2> attrs { start, result };
        std::array<uint8_t, 2> tupleCounts { 16, 4 };
        std::array<uint8_t, 2> arrayDepths { startType.arrayDepth, startType.arrayDepth };
        std::array<AttributeRole, 2> rolesBuf { AttributeRole::eNone, AttributeRole::eQuaternion };
        node.iNode->resolvePartiallyCoupledAttributes(node, attrs.data(), tupleCounts.data(), 
                                                    arrayDepths.data(), rolesBuf.data(), attrs.size());
    }
};

REGISTER_OGN_NODE()

}
}
}
