using System.Collections.Generic;
using System.Diagnostics;
using Unity.InferenceEngine;
using Unity.MLAgents.Actuators;
using Unity.MLAgents.Inference;
using Unity.MLAgents.Policies;
using Unity.MLAgents.Sensors;
using UnityEngine;

#if MLA_UNITY_ANALYTICS_MODULE && ENABLE_CLOUD_SERVICES_ANALYTICS
using UnityEngine.Analytics;
#endif


#if UNITY_EDITOR
using UnityEditor;
#if MLA_UNITY_ANALYTICS_MODULE
using UnityEditor.Analytics;
#endif // MLA_UNITY_ANALYTICS_MODULE
#endif // UNITY_EDITOR


namespace Unity.MLAgents.Analytics
{
    internal class InferenceAnalytics
    {


#if UNITY_EDITOR && MLA_UNITY_ANALYTICS_MODULE && ENABLE_CLOUD_SERVICES_ANALYTICS
        /// <summary>
        /// Models that we've already sent events for.
        /// </summary>
        private static HashSet<ModelAsset> s_SentModels;
#endif

        static bool EnableAnalytics()
        {
#if UNITY_EDITOR && MLA_UNITY_ANALYTICS_MODULE && ENABLE_CLOUD_SERVICES_ANALYTICS

            if (s_SentModels == null)
            {
                s_SentModels = new HashSet<ModelAsset>();
            }

            return true;

#else  // no editor, no analytics
            return false;
#endif
        }

        public static bool IsAnalyticsEnabled()
        {
#if UNITY_EDITOR
            return EditorAnalytics.enabled;
#else
            return false;
#endif
        }

        /// <summary>
        /// Send an analytics event for the NNModel when it is set up for inference.
        /// No events will be sent if analytics are disabled, and at most one event
        /// will be sent per model instance.
        /// </summary>
        /// <param name="nnModel">The NNModel being used for inference.</param>
        /// <param name="behaviorName">The BehaviorName of the Agent using the model</param>
        /// <param name="inferenceDevice">Whether inference is being performed on the CPU or GPU</param>
        /// <param name="sensors">List of ISensors for the Agent. Used to generate information about the observation space.</param>
        /// <param name="actionSpec">ActionSpec for the Agent. Used to generate information about the action space.</param>
        /// <param name="actuators">List of IActuators for the Agent. Used to generate information about the action space.</param>
        [Conditional("MLA_UNITY_ANALYTICS_MODULE")]
        public static void InferenceModelSet(
            ModelAsset nnModel,
            string behaviorName,
            InferenceDevice inferenceDevice,
            IList<ISensor> sensors,
            ActionSpec actionSpec,
            IList<IActuator> actuators
        )
        {
#if UNITY_EDITOR && MLA_UNITY_ANALYTICS_MODULE && ENABLE_CLOUD_SERVICES_ANALYTICS
            // The event shouldn't be able to report if this is disabled but if we know we're not going to report
            // Lets early out and not waste time gathering all the data
            if (!IsAnalyticsEnabled())
                return;

            if (!EnableAnalytics())
                return;

            var added = s_SentModels.Add(nnModel);

            if (!added)
            {
                // We previously added this model. Exit so we don't resend.
                return;
            }

            var data = GetEventForModel(nnModel, behaviorName, inferenceDevice, sensors, actionSpec, actuators);
            // Note - to debug, use JsonUtility.ToJson on the event.
            // Debug.Log(JsonUtility.ToJson(data, true));
            if (AnalyticsUtils.s_SendEditorAnalytics)
            {
                EditorAnalytics.SendAnalytic(data);
            }
#endif
        }

        /// <summary>
        /// Generate an InferenceEvent for the model.
        /// </summary>
        /// <param name="nnModel"></param>
        /// <param name="behaviorName"></param>
        /// <param name="inferenceDevice"></param>
        /// <param name="sensors"></param>
        /// <param name="actionSpec"></param>
        /// <param name="actuators"></param>
        /// <returns>`InferenceEvent` from the input model.</returns>
        internal static InferenceEvent GetEventForModel(
            ModelAsset nnModel,
            string behaviorName,
            InferenceDevice inferenceDevice,
            IList<ISensor> sensors,
            ActionSpec actionSpec,
            IList<IActuator> actuators
        )
        {
            var sentisModel = ModelLoader.Load(nnModel);
            using var sentisModelInfo = new SentisModelInfo(sentisModel);
            var inferenceEvent = new InferenceEvent();

            // Hash the behavior name so that there's no concern about PII or "secret" data being leaked.
            inferenceEvent.BehaviorName = AnalyticsUtils.Hash(AnalyticsConstants.k_VendorKey, behaviorName);

            inferenceEvent.SentisModelVersion = sentisModelInfo.Version;
            inferenceEvent.SentisModelProducer = sentisModel.ProducerName;
            inferenceEvent.MemorySize = sentisModelInfo.MemorySize;
            inferenceEvent.InferenceDevice = (int)inferenceDevice;

            // TODO deprecate tensorflow conversion
            if (sentisModel.ProducerName == "Script")
            {
                // .nn files don't have these fields set correctly. Assign some placeholder values.
                inferenceEvent.SentisModelSource = "NN";
                inferenceEvent.SentisModelProducer = "tensorflow_to_barracuda.py";
            }

#if UNITY_EDITOR
            var sentisPackageInfo = UnityEditor.PackageManager.PackageInfo.FindForAssembly(typeof(Tensor).Assembly);
            inferenceEvent.SentisPackageVersion = sentisPackageInfo.version;
#else
            inferenceEvent.SentisPackageVersion = null;
#endif

            inferenceEvent.ActionSpec = EventActionSpec.FromActionSpec(actionSpec);
            inferenceEvent.ObservationSpecs = new List<EventObservationSpec>(sensors.Count);
            foreach (var sensor in sensors)
            {
                inferenceEvent.ObservationSpecs.Add(EventObservationSpec.FromSensor(sensor));
            }

            inferenceEvent.ActuatorInfos = new List<EventActuatorInfo>(actuators.Count);
            foreach (var actuator in actuators)
            {
                inferenceEvent.ActuatorInfos.Add(EventActuatorInfo.FromActuator(actuator));
            }

            inferenceEvent.TotalWeightSizeBytes = GetModelWeightSize(sentisModel);
            inferenceEvent.ModelHash = GetModelHash(sentisModel);
            return inferenceEvent;
        }

        /// <summary>
        /// Compute the total model weight size in bytes.
        /// This corresponds to the "Total weight size" display in the Sentis inspector,
        /// and the calculations are the same.
        /// </summary>
        /// <param name="sentisModel"></param>
        /// <returns>The total model weight size in bytes.</returns>
        static long GetModelWeightSize(Model sentisModel)
        {
            long totalWeightsSizeInBytes = 0;
            for (var c = 0; c < sentisModel.constants.Count; c++)
            {
                totalWeightsSizeInBytes += sentisModel.constants[c].lengthBytes;
            }
            return totalWeightsSizeInBytes;
        }

        /// <summary>
        /// Wrapper around Hash128 that supports Append(float[], int, int)
        /// </summary>
        struct MLAgentsHash128
        {
            private Hash128 m_Hash;

            public void Append(float[] values, int count)
            {
                if (values == null)
                {
                    return;
                }

                // Pre-2020 versions of Unity don't have Hash128.Append() (can only hash strings and scalars)
                // For these versions, we'll hash element by element.
#if UNITY_2020_1_OR_NEWER
                m_Hash.Append(values, 0, count);
#else
                for (var i = 0; i < count; i++)
                {
                    var tempHash = new Hash128();
                    HashUtilities.ComputeHash128(ref values[i], ref tempHash);
                    HashUtilities.AppendHash(ref tempHash, ref m_Hash);
                }
#endif
            }

            public void Append(string value)
            {
                var tempHash = Hash128.Compute(value);
                HashUtilities.AppendHash(ref tempHash, ref m_Hash);
            }

            public override string ToString()
            {
                return m_Hash.ToString();
            }
        }

        /// <summary>
        /// Compute a hash of the model's layer data and return it as a string.
        /// A subset of the layer weights are used for performance.
        /// This increases the chance of a collision, but this should still be extremely rare.
        /// </summary>
        /// <param name="sentisModel"></param>
        /// <returns>The hash of the model's layer data.</returns>
        static string GetModelHash(Model sentisModel)
        {
            var hash = new MLAgentsHash128();

            foreach (var constant in sentisModel.constants)
            {
                hash.Append(constant.ToString());
            }

            return hash.ToString();
        }
    }
}