// Copyright (c) 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 <OgnTutorialExtendedTypesDatabase.h>
#include <algorithm>

//
// Attributes whose data types resolve at runtime ("any" or "union" types) are resolved by having connections made
// to them of a resolved type. Say you have a chain of A->B->C where B has inputs and outputs of these types. The
// connection from A->B will determine the type of data at B's input and the connection B->C will determine the type
// of data at B's output (assuming A's outputs and C's inputs are well-defined types).
//
// For this reason it is the node's responsibility to verify the type resolution of the attributes as part of the
// compute method. Any unresolved types (db.Xputs.attrName().resolved() == false) that are required by the compute
// should result in a warning and compute failure. Any attributes resolved to incompatible types, for example an input
// that resolves to a string where a number is needed, should also result in a warning and compute failure.
//
// It is up to the node to decide how flexible the resolution requirements are to be. In the string/number case above
// the node may choose to parse the string as a number instead of failing, or using the length of the string as the
// input number. The only requirement from OmniGraph is that the node handle all of the resolution types it has
// claimed it will handle in the .ogn file. "any" attributes must handle all data types, even if some types result in
// warnings or errors. "union" attributes must handle all types specified in the union.
//

class OgnTutorialExtendedTypes
{
public:
    static bool compute(OgnTutorialExtendedTypesDatabase& db)
    {
        bool computedOne = false;

        auto typeWarning = [&](const char* message, const Type& type1, const Type& type2) {
            db.logWarning("%s (%s -> %s)", message, getOgnTypeName(type1).c_str(), getOgnTypeName(type2).c_str());
        };
        auto typeError = [&](const char* message, const Type& type1, const Type& type2) {
            db.logError("%s (%s -> %s)", message, getOgnTypeName(type1).c_str(), getOgnTypeName(type2).c_str());
        };

        auto computeSimpleValues = [&]() {
            // ====================================================================================================
            // Compute for the union types that resolve to simple values. 
            // Accepted value types are floats and tokens. As these were the only types specified in the union definition
            // the node does not have to worry about other numeric types, such as int or double.

            // The node can decide what the meaning of an attempt to compute with unresolved types is.
            // For this particular node they are treated as silent success.
            const auto& floatOrToken = db.inputs.floatOrToken();
            auto& doubledResult = db.outputs.doubledResult();

            if (floatOrToken.resolved() && doubledResult.resolved())
            {
                // Check for an exact type match for the input and output
                if (floatOrToken.type() != doubledResult.type())
                {
                    // Mismatched types are possible, and result in no compute
                    typeWarning("Simple resolved types do not match", floatOrToken.type(), doubledResult.type());
                    return false;
                }

                // When extracting extended types the templated get<> method returns an object that contains the cast data.
                // It can be cast to a boolean for quick checks for matching types.
                //
                // Note: The single "=" in these if statements is intentional. It facilitates one-line set-and-test of the
                //       typed values.
                //
                if (auto floatValue = floatOrToken.get<float>())
                {
                    // Once the existence of the cast type is verified it can be dereferenced to get at the raw data,
                    // whose types are described in the tutorial on bundled data.
                    if (auto doubledValue = doubledResult.get<float>())
                    {
                        *doubledValue = *floatValue * 2.0f;
                    }
                    else
                    {
                        // This could be an assert because it should never happen. The types were confirmed above to match,
                        // so they should have cast to the same types without incident.
                        typeError("Simple types were matched as bool then failed to cast properly", floatOrToken.type(), doubledResult.type());
                        return false;
                    }
                }
                else if (auto tokenValue = floatOrToken.get<OgnToken>())
                {
                    if (auto doubledValue = doubledResult.get<OgnToken>())
                    {
                        std::string inputString{ db.tokenToString(*tokenValue) };
                        inputString += inputString;
                        *doubledValue = db.stringToToken(inputString.c_str());
                    }
                    else
                    {
                        // This could be an assert because it should never happen. The types were confirmed above to match,
                        // so they should have cast to the same types without incident.
                        typeError("Simple types were matched as token then failed to cast properly", floatOrToken.type(), doubledResult.type());
                        return false;
                    }
                }
                else
                {
                    // As Union types are supposed to restrict the data types being passed in to the declared types
                    // any unrecognized types are an error, not a warning.
                    typeError("Simple types resolved to unknown types", floatOrToken.type(), doubledResult.type());
                    return false;
                }
            }
            else
            {
                // Unresolved types are reasonable, resulting in no compute
                return true;
            }
            return true;
        };
    
        auto computeArrayValues = [&]() {
            // ====================================================================================================
            // Compute for the union types that resolve to arrays.
            // Accepted value types are arrays of bool or arrays of float, which are extracted as interfaces to
            // those values so that resizing can happen transparently through the fabric.
            //
            // These interfaces are similar to what you've seen in regular array attributes - they support resize(),
            // operator[], and range-based for loops.
            //
            const auto& toNegate = db.inputs.toNegate();
            auto& negatedResult = db.outputs.negatedResult();

            if (toNegate.resolved() && negatedResult.resolved())
            {
                // Check for an exact type match for the input and output
                if (toNegate.type() != negatedResult.type())
                {
                    // Mismatched types are possible, and result in no compute
                    typeWarning("Array resolved types do not match", toNegate.type(), negatedResult.type());
                    return false;
                }

                // Extended types can be any legal attribute type. Here the types in the extended attribute can be either
                // an array of booleans or an array of integers.
                if (auto boolArray = toNegate.get<bool[]>())
                {
                    auto valueAsBoolArray = negatedResult.get<bool[]>();
                    if (valueAsBoolArray)
                    {
                        valueAsBoolArray.resize( boolArray->size() );
                        size_t index{ 0 };
                        for (auto& value : *boolArray)
                        {
                            (*valueAsBoolArray)[index++] = ! value;
                        }
                    }
                    else
                    {
                        // This could be an assert because it should never happen. The types were confirmed above to match,
                        // so they should have cast to the same types without incident.
                        typeError("Array types were matched as bool[] then failed to cast properly", toNegate.type(), negatedResult.type());
                        return false;
                    }
                }
                else if (auto floatArray = toNegate.get<float[]>())
                {
                    auto valueAsFloatArray = negatedResult.get<float[]>();
                    if (valueAsFloatArray)
                    {
                        valueAsFloatArray.resize( floatArray->size() );
                        size_t index{ 0 };
                        for (auto& value : *floatArray)
                        {
                            (*valueAsFloatArray)[index++] = - value;
                        }
                    }
                    else
                    {
                        // This could be an assert because it should never happen. The types were confirmed above to match,
                        // so they should have cast to the same types without incident.
                        typeError("Array types were matched as float[] then failed to cast properly", toNegate.type(), negatedResult.type());
                        return false;
                    }
                }
                else
                {
                    // As Union types are supposed to restrict the data types being passed in to the declared types
                    // any unrecognized types are an error, not a warning.
                    typeError("Array type not recognized", toNegate.type(), negatedResult.type());
                    return false;
                }
            }
            else
            {
                // Unresolved types are reasonable, resulting in no compute
                return true;
            }
            return true;
        };

        auto computeTupleValues = [&]() {
            // ====================================================================================================
            // Compute for the "any" types that only handle tuple values.  In practice you'd only use "any" when the
            // type of data you handle is unrestricted. This is more an illustration to show how in practical use the
            // two types of attribute are accessed exactly the same way, the only difference is restrictions that the
            // OmniGraph system will put on potential connections.
            //
            // For simplicity this node will treat unrecognized type as a warning with success.
            // Full commentary and error checking is elided as it will be the same as for the above examples.
            // The algorithm for tuple values is a component-wise negation.
            const auto& tupleInput = db.inputs.tuple();
            auto& tupleOutput = db.outputs.tuple();

            if (tupleInput.resolved() && tupleOutput.resolved())
            {
                if (tupleInput.type() != tupleOutput.type())
                {
                    typeWarning("Tuple resolved types do not match", tupleInput.type(), tupleOutput.type());
                    return false;
                }

                // This node will only recognize the float[3] and int[2] cases, to illustrate that tuple count and
                // base type are both flexible.
                if (auto float3Input = tupleInput.get<float[3]>())
                {
                    if (auto float3Output = tupleOutput.get<float[3]>())
                    {
                        (*float3Output)[0] = -(*float3Input)[0];
                        (*float3Output)[1] = -(*float3Input)[1];
                        (*float3Output)[2] = -(*float3Input)[2];
                    }
                }
                else if (auto int2Input = tupleInput.get<int[2]>())
                {
                    if (auto int2Output = tupleOutput.get<int[2]>())
                    {
                        (*int2Output)[0] = -(*int2Input)[0];
                        (*int2Output)[1] = -(*int2Input)[1];
                    }
                }
                else
                {
                    // As "any" types are not restricted in their data types but this node is only handling two of
                    // them an unrecognized type is just unimplemented code.
                    typeWarning("Unimplemented type combination", tupleInput.type(), tupleOutput.type());
                    return true;
                }
            }
            else
            {
                // Unresolved types are reasonable, resulting in no compute
                return true;
            }
            return true;
        };
    
        auto computeFlexibleValues = [&]() {
            // ====================================================================================================
            // Complex union type that handles both simple values and an array of tuples. It illustrates how the
            // data types in a union do not have to be related in any way.
            //
            // Full commentary and error checking is elided as it will be the same as for the above examples.
            // The algorithm for tuple array values is to negate everything in the float3 array values, and to reverse
            // the string for string values.
            const auto& flexibleInput = db.inputs.flexible();
            auto& flexibleOutput = db.outputs.flexible();

            if (flexibleInput.resolved() && flexibleOutput.resolved())
            {
                if (flexibleInput.type() != flexibleOutput.type())
                {
                    typeWarning("Flexible resolved types do not match", flexibleInput.type(), flexibleOutput.type());
                    return false;
                }

                // Arrays of tuples are handled with the same interface as with normal attributes.
                if (auto float3ArrayInput = flexibleInput.get<float[][3]>())
                {
                    if (auto float3ArrayOutput = flexibleOutput.get<float[][3]>())
                    {
                        float3ArrayOutput.resize( float3ArrayInput.size() );
                        size_t index{ 0 };
                        for (auto& value : *float3ArrayInput)
                        {
                            (*float3ArrayOutput)[index][0] = - value[0];
                            (*float3ArrayOutput)[index][1] = - value[1];
                            (*float3ArrayOutput)[index][2] = - value[2];
                            index++;
                        }
                    }
                }
                else if (auto tokenInput = flexibleInput.get<OgnToken>())
                {
                    if (auto tokenOutput = flexibleOutput.get<OgnToken>())
                    {
                        std::string toReverse{ db.tokenToString(*tokenInput) };
                        std::reverse( toReverse.begin(), toReverse.end() );
                        *tokenOutput = db.stringToToken(toReverse.c_str());
                    }
                }
                else
                {
                    typeError("Unrecognized type combination", flexibleInput.type(), flexibleOutput.type());
                    return false;
                }
            }
            else
            {
                // Unresolved types are reasonable, resulting in no compute
                return true;
            }

            return true;
        };

        // This approach lets either section fail while still computing the other.
        computedOne = computeSimpleValues();
        computedOne = computeArrayValues() || computedOne;
        computedOne = computeTupleValues() || computedOne;
        computedOne = computeFlexibleValues() || computedOne;

        if (! computedOne)
        {
            db.logWarning("None of the inputs had resolved type, resulting in no compute");
        }
        return ! computedOne;
    }

    static void onConnectionTypeResolve(const NodeObj& nodeObj)
    {
        // The attribute types resolve in pairs
        AttributeObj pairs[][2] {
            {
                nodeObj.iNode->getAttributeByToken(nodeObj, inputs::floatOrToken.token()),
                nodeObj.iNode->getAttributeByToken(nodeObj, outputs::doubledResult.token())
            },
            {
                nodeObj.iNode->getAttributeByToken(nodeObj, inputs::toNegate.token()),
                nodeObj.iNode->getAttributeByToken(nodeObj, outputs::negatedResult.token())
            },
            {
                nodeObj.iNode->getAttributeByToken(nodeObj, inputs::tuple.token()),
                nodeObj.iNode->getAttributeByToken(nodeObj, outputs::tuple.token())
            },
            {
                nodeObj.iNode->getAttributeByToken(nodeObj, inputs::flexible.token()),
                nodeObj.iNode->getAttributeByToken(nodeObj, outputs::flexible.token())
            }
        };
        for (auto& pair : pairs)
        {
            nodeObj.iNode->resolveCoupledAttributes(nodeObj, &pair[0], 2);
        }
    }
};

REGISTER_OGN_NODE()
