//
// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE.md file in the project root for full license information.
//

#include "UserDefinedFunction.h"
#include "UserFunctionFactory.h"
#include "Serialization.h"
#include "PrimitiveFunction.h"
#include "CompositeFunction.h"


namespace CNTK
{

    static Internal::UDFDeserializeCallbackWrapperPtr s_SWIGCallbackWrapper;

    void Internal::RegisterUDFDeserializeCallbackWrapper(UDFDeserializeCallbackWrapperPtr callbackPtr) 
    {
        s_SWIGCallbackWrapper = callbackPtr;
    }

    static const std::wstring s_nativeUDFTypeValue = L"NativeUserDefinedFunction" ;

    static std::unordered_map<std::wstring, std::pair<std::wstring, std::wstring>> s_deserializedUDFsRegistry;

    Function::Function(const std::vector<Variable>& inputs, const std::wstring& name)
        : Function(inputs, {}, name)
    {}

    Function::Function(const std::vector<Variable>& inputs, const Dictionary& functionConfig, const std::wstring& name)
        : Function(inputs, functionConfig, nullptr, name, Internal::GenerateUid(L"UserDefinedFunction"))
    {}

    /*static*/ FunctionPtr Function::DeserializeNativeImpl(const std::vector<Variable>& inputs, const std::wstring& name,  const Dictionary& dict)
    {
        static const vector<std::wstring> s_requiredDictionaryKeys = { userDefinedStateKey, udfModuleNameKey, udfFactoryMethodNameKey, opKey };
        ValidateDictionary<PrimitiveFunction>(dict, s_requiredDictionaryKeys, s_nativeUDFTypeValue, s_serializationVersion);

        auto state = dict[userDefinedStateKey].Value<Dictionary>();
        auto opName = dict[opKey].Value<wstring>();
        auto moduleName = dict[udfModuleNameKey].Value<wstring>();
        auto methodName = dict[udfFactoryMethodNameKey].Value<wstring>();

        FunctionPtr udf = nullptr;

        auto callback = Function::GetUDFDeserializeCallback(opName);
        if (callback != nullptr) 
        {
            udf = callback->operator()(inputs, name, state);
        }
        else 
        {
            Microsoft::MSR::CNTK::Plugin plugin;
            auto factoryMethod = (UserFunctionFactoryMethodType)(plugin.Load(moduleName, msra::strfun::utf8(methodName), /*isCNTKPlugin =*/ false));
            udf = FunctionPtr(factoryMethod(inputs.data(), inputs.size(), &state, name.c_str()));
        }

        if (udf == nullptr) 
        {
            RuntimeError("Unable to reconstruct the native UserFunction with op name '%S'", opName.c_str());
        }

        s_deserializedUDFsRegistry[opName] = { moduleName, methodName };
        
        return udf;
    }

    /*virtual*/ std::wstring Function::ModuleName() const
    {
        auto it = s_deserializedUDFsRegistry.find(OpName());
        if (it != s_deserializedUDFsRegistry.end())
        {
            auto moduleAndMethodPair = it->second;
            return moduleAndMethodPair.first;
        }
        
        // this op name was never registered in the s_deserializedUDFsRegistry (which only happens during the deserialization),
        // then use user factory as a fallback (this udf must have been registed, so that an instance could be created).
        return s_userFunctionFactory->GetModuleName(OpName());
    }

    /*virtual*/ std::wstring Function::DeserializeMethodName() const
    {
        auto it = s_deserializedUDFsRegistry.find(OpName());
        if (it != s_deserializedUDFsRegistry.end())
        {
            auto moduleAndMethodPair = it->second;
            return moduleAndMethodPair.second;
        }

        // this op name was never registered in the s_deserializedUDFsRegistry (which only happens during the deserialization),
        // then use user factory as a fallback (this udf must have been registed, so that an instance could be created).
        return Function::s_userFunctionFactory->GetFactoryMethodName(OpName());
    }

    Dictionary Function::SerializeNativeImpl() const
    {
        Dictionary dict;
        dict[userDefinedStateKey] = Serialize();
        dict[udfModuleNameKey] = ModuleName();
        dict[udfFactoryMethodNameKey] = DeserializeMethodName();
        dict[opKey] = OpName();
        dict[versionKey] = s_serializationVersion;
        dict[typeKey] = s_nativeUDFTypeValue;
        return dict;
    }

    static const std::wstring s_userDefinedFunctionTypeValue = L"UserDefinedFunction";

    /*static*/ bool UDFUtils::IsUDF(const FunctionPtr& f)
    {
        return (dynamic_cast<const PrimitiveFunction*>(f.get()) == nullptr) &&
            (dynamic_cast<const CompositeFunction*>(f.get()) == nullptr);
    }

    /*static*/ bool UDFUtils::IsUDF(const Dictionary& dict)
    {
        return (dict.Contains(typeKey) && dict[typeKey].Value<std::wstring>() == s_userDefinedFunctionTypeValue);
    }

    /*static*/ bool UDFUtils::IsNativeUDF(const Dictionary& dict)
    {
        assert(IsUDF(dict));
        return (dict.Contains(nativeUDFKey) && dict[nativeUDFKey].Value<bool>() == true);
    }

    /*static*/ Dictionary UDFUtils::Serialize(const FunctionPtr& f)
    {
        Dictionary dict = SerializeCommonFunctionAttributes(*f, s_serializationVersion, s_userDefinedFunctionTypeValue);
        bool native = f->IsNative();
        dict[nativeUDFKey] = native;
        dict[userDefinedStateKey] = (native) ? f->SerializeNativeImpl() : f->Serialize();
        return dict;
    }

    /*static*/ FunctionPtr UDFUtils::Deserialize(const Dictionary& dict,
        const unordered_map<std::wstring, Variable>& uidToVariableMap,
        const DeviceDescriptor& device)
    {
        static const vector<std::wstring> s_requiredDictionaryKeys = { typeKey, uidKey, inputsKey, userDefinedStateKey };
        ValidateDictionary<PrimitiveFunction>(dict, s_requiredDictionaryKeys, s_userDefinedFunctionTypeValue, s_serializationVersion);

        const auto& uid = dict[uidKey].Value<std::wstring>();
        std::wstring name = L"";
        if (dict.Contains(nameKey))
            name = dict[nameKey].Value<std::wstring>();

        auto inputs = GetInputVariables(dict, uidToVariableMap, s_serializationVersion);

        auto state = dict[userDefinedStateKey].Value<Dictionary>();

        FunctionPtr udf;

        if (IsNativeUDF(dict))
        {
            udf = Function::DeserializeNativeImpl(inputs, name, state);
        }
        else if (s_SWIGCallbackWrapper != nullptr)
        {
            // If we're being called from SWIG, the actual deserializer should be registered by
            // the target language CNTK implementation (i.e., cntk_py for Python)
            udf = s_SWIGCallbackWrapper->operator()(inputs, name, state);
        }

        if (udf == nullptr)
        {
            RuntimeError("Unable to reconstruct a user-defined function (name = %S, uid = %S). "
                "Please make sure to specify a valid UDF deserializer.", name.c_str(), uid.c_str());
        }

        // Restore the original uid, which other functions in the graph depend on
        // (their inputs refer to the uids of this UDF outputs, which are generated base on the uid of this UDF).
        udf->m_uid = uid;

        return udf;
    }
}