/*=========================================================================

   Library: iMSTK-Unity

   Copyright (c) Kitware, Inc. 

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0.txt

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.

=========================================================================*/

using System.Collections.Generic;
using System.Linq;
using UnityEngine;
using Unity.Profiling;

namespace ImstkUnity
{
    [System.Serializable]
    public class  PbdModelConfiguration
    {
        public Vector3 gravity = new Vector3(0.0f, -9.81f, 0.0f);
        public int iterations = 10;
        public bool useRealtime = true;
        public double dt = 0.01;

        public double linearDampingCoeff = 0.01;
        public double angularDampingCoeff = 0.01;

        public bool doPartitioning = false;

        public static PbdModelConfiguration Default()
        {
            PbdModelConfiguration result = new PbdModelConfiguration();
            var defaultConfig = new Imstk.PbdModelConfig();
            Imstk.Vec3d gravity = defaultConfig.m_gravity;
            result.gravity = gravity.ToUnityVec();
            result.iterations = (int)defaultConfig.m_iterations;
            result.useRealtime = true;
            result.dt = defaultConfig.m_dt;
            result.linearDampingCoeff = defaultConfig.m_linearDampingCoeff;
            result.angularDampingCoeff = defaultConfig.m_angularDampingCoeff;
            result.doPartitioning = defaultConfig.m_doPartitioning;

            return result;
        }

        public PbdModelConfiguration DeepCopy()
        {
            PbdModelConfiguration result = new PbdModelConfiguration();
            result.gravity = gravity;
            result.iterations = iterations;
            result.useRealtime = useRealtime;
            result.dt = dt;
            result.linearDampingCoeff = linearDampingCoeff;
            result.angularDampingCoeff = angularDampingCoeff;
            result.doPartitioning = doPartitioning;

            return result;
        }

    }

    [AddComponentMenu("iMSTK/SimulationManager")]
    [DefaultExecutionOrder(-99)]
    public class SimulationManager : MonoBehaviour
    {
        public static Imstk.SceneManager sceneManager = null;

        static readonly ProfilerMarker s_AdvancePerfMarker = new ProfilerMarker("SimulationManager.Advance");

        // This overwrites the Project Settings->Fixed Timestep as that one is per project not scene
        // This could be a point of pain when integrating with other Unity plugins
        public float sceneFixedTimestep = 0.01f;

        public bool rbdUseRealtime = false;
        public Vector3 rigidBodyGravity = new Vector3(0.0f, -9.81f, 0.0f);
        public int rigidBodyMaxNumIterations = 11;
        public double rigidBodyDt = 0.01;

        public bool writeTaskGraph = false;

        private Imstk.CacheOutput output;


        public PbdModelConfiguration pbdModelConfiguration = new PbdModelConfiguration();
        public static Imstk.RigidBodyModel2 rigidBodyModel = null;
        public static Imstk.PbdModel pbdModel = null;

        /// <summary>
        /// Returns all components in the scene of a given type.
        /// This allows us to collect the iMSTK components and insert them into iMSKT no
        /// matter where they are located. Components that are not "enable" will be disregarded
        /// </summary>
        /// <typeparam name="T">Class of the Component that you are looking for, needs to be a Monobehavior</typeparam>
        /// <returns>List of all _active_ components in the scene of type T</returns>
        public static List<T> GetAllComponents<T>() where T : MonoBehaviour 
        {
            List<T> behaviours = new List<T>();
            List<GameObject> objects = FindObjectsOfType<GameObject>().ToList();
            foreach (GameObject obj in objects)
            {
                foreach (var behaviour in obj.GetComponents<T>())
                {
                    if (behaviour.enabled)
                    {
                        behaviours.Add(behaviour);
                    }
                }
            }
            return behaviours;
        }

        private void CreateRigidBodyModel()
        {
            rigidBodyModel = new Imstk.RigidBodyModel2();
            Imstk.RigidBodyModel2Config rbdConfig = rigidBodyModel.getConfig();
            rbdConfig.m_gravity = rigidBodyGravity.ToImstkVec();
            rbdConfig.m_maxNumIterations = (uint)rigidBodyMaxNumIterations;
            rbdConfig.m_dt = rigidBodyDt;
        }

        private void Awake()
        {
            // Get the settings
            ImstkSettings settings = ImstkSettings.Instance();
            //if (settings.useOptimalNumberOfThreads)
            //    settings.numThreads = 0;

            Imstk.Logger.startLogger();
            output = Imstk.Logger.instance().getCacheOutput();

            // Create the simulation manager
            sceneManager = new Imstk.SceneManager();
            sceneManager.setActiveScene(new Imstk.Scene("DefaultScene"));
            sceneManager.getActiveScene().getConfig().writeTaskGraph = writeTaskGraph;

            // Create A single RigidBodyModel to share for all rigid bodies used in the scene
            CreateRigidBodyModel();

            // Create A single PbdModel to share for all pbd bodies used in the scene 
            CreatePbdModel();

            // Override the Unity fixed delta time
            Time.fixedDeltaTime = sceneFixedTimestep;

        }

        private void Start()
        {

            // It seems that InitManager needs to come AFTER the call that creates the
            // device inside of the OpenHapticsDevice
            IntializeImstkStructures();

#if IMSTK_USE_OPENHAPTICS
            OpenHapticsDevice.InitManager();
#endif

            sceneManager.init();

            // Start order
            {
                List<ImstkBehaviour> behaviours = GetAllComponents<ImstkBehaviour>();
                foreach (ImstkBehaviour behaviour in behaviours)
                {
                    behaviour.ImstkStart();
                }
            }

            // #refactor should follow the same pattern as all
            // i.e. Get all Managers in the scene and start them
            // use same pattern for both managers
#if IMSTK_USE_VRPN
            if (VrpnDeviceManager.Instance != null) VrpnDeviceManager.Instance.StartManager();
#endif
#if IMSTK_USE_OPENHAPTICS
            OpenHapticsDevice.StartManager();
#endif
        }

        // Update the imstk scene within the fixed update of Unity
        public void FixedUpdate()
        {
            if (rbdUseRealtime)
            {
                rigidBodyModel.getConfig().m_dt = Time.fixedDeltaTime;
            }

            if (pbdModelConfiguration.useRealtime)
            {
                pbdModel.getConfig().m_dt = Time.fixedDeltaTime;
            }
            s_AdvancePerfMarker.Begin();
            sceneManager.getActiveScene().advance(Time.fixedDeltaTime);
            s_AdvancePerfMarker.End();
        }

        public void OnApplicationQuit()
        {
            // Destroy order
            {
                List<ImstkBehaviour> behaviours = GetAllComponents<ImstkBehaviour>();
                foreach (ImstkBehaviour behaviour in behaviours)
                {
                    behaviour.ImstkDestroy();
                }
            }

            sceneManager.uninit();

#if IMSTK_USE_OPENHAPTICS
            OpenHapticsDevice.StopManager(); // Stops in async
#endif
#if IMSTK_USE_VRPN
            if (VrpnDeviceManager.Instance != null) VrpnDeviceManager.Instance.StopManager();
#endif

            // These are static so we need to make sure to set them to null on quit
            sceneManager = null;
            rigidBodyModel = null;
        }

        private void Update()
        {
            LogToUnity();
        }

        private void IntializeImstkStructures()
        {
            // SimulationManager initializes objects in a particular order
            List<ImstkBehaviour> behaviours = GetAllComponents<ImstkBehaviour>();
            foreach (ImstkBehaviour behaviour in behaviours)
            {
                behaviour.ImstkInit(); 
                
            }

            // We need to ensure all objects are created before interactions and controllers
            // are setup using them
            List<ImstkInteractionBehaviour> interactions = GetAllComponents<ImstkInteractionBehaviour>();
            foreach (ImstkInteractionBehaviour behaviour in interactions)
            {
                Imstk.SceneObject interaction = behaviour.GetImstkInteraction();
                if (interaction != null)
                {
                    sceneManager.getActiveScene().addInteraction(interaction);
                }
            }

            List<ImstkControllerBehaviour> controllers = GetAllComponents<ImstkControllerBehaviour>();
            foreach (ImstkControllerBehaviour behaviour in controllers)
            {
                // Currently only support tracking device controls
                Imstk.TrackingDeviceControl control =
                    behaviour.GetController() as Imstk.TrackingDeviceControl;
                if (control != null)
                {
                    sceneManager.getActiveScene().addControl(control);
                }
            }
        }

        private void LogToUnity()
        {
            while (output.hasMessages())
                Debug.Log(output.popLastMessage());
        }

        private void CreatePbdModel()
        {
            pbdModel = new Imstk.PbdModel();
            Imstk.PbdModelConfig config = new Imstk.PbdModelConfig();

            config.m_dt = pbdModelConfiguration.dt;
            config.m_gravity = pbdModelConfiguration.gravity.ToImstkVec();
            config.m_iterations = (uint)pbdModelConfiguration.iterations;
            config.m_linearDampingCoeff = pbdModelConfiguration.linearDampingCoeff;
            config.m_angularDampingCoeff = pbdModelConfiguration.angularDampingCoeff;
            config.m_doPartitioning = pbdModelConfiguration.doPartitioning;
            pbdModel.configure(config);
            pbdModel.setTimeStepSizeType((pbdModelConfiguration.useRealtime) ? Imstk.TimeSteppingType.RealTime : Imstk.TimeSteppingType.Fixed);
        }

    }
}