// // /*===============================================================================
// // Copyright (C) 2020 PhantomsXR Ltd. All Rights Reserved.
// //
// // This file is part of the AR-MOD SDK.
// //
// // The AR-MOD SDK cannot be copied, distributed, or made available to
// // third-parties for commercial purposes without written permission of PhantomsXR Ltd.
// //
// // Contact info@phantomsxr.com for licensing requests.
// // ===============================================================================*/

using Unity.Barracuda;
using Unity.VisualScripting;

namespace com.PhantomsXR.VisualScriptingExtension.Runtime
{
    [UnitCategory("ARMOD/NNEInference/Tensor Extension"), UnitTitle("ArgMax"), UnitShortTitle("ArgMax")]
    public class ArgMax : ARMODBaseWithResultUnit
    {
        [DoNotSerialize] public ValueInput tensor;
        private int[] result;

        protected override void Definition()
        {
            base.Definition();
            tensor = ValueInput(nameof(Tensor), (Tensor) null);
            Result = ValueOutput("", _flow => result);
            Requirement(tensor, inputTrigger);
        }

        public override ControlOutput Execute(Flow _arg)
        {
            var tmp_Tensor = _arg.GetValue<Tensor>(tensor);
            result = tmp_Tensor.ArgMax();
            return outputTrigger;
        }
    }
}