using System;
using System.Collections.Generic;
using UnityEngine;
#if UNITY_MATHEMATICS
using Unity.Mathematics;
using static Unity.Mathematics.math;
#endif
using NanoBrain.Unity;

namespace NanoBrain {

    /// <summary>
    /// A Cluster combines a collection of Nuclei to implement reusable behaviour
    /// </summary>
    /// A Cluster is an instantiation of a ClusterPrefab.
    /// Clusters can be nested inside other clusters.
    [Serializable]
    public class Cluster : Nucleus {
        // It may be that clusters will not be nuclei anymore in the future....

        /// <summary>
        /// The prefab used to create this cluster
        /// </summary>
        /// Cluster should always be created from prefabs
        public ClusterPrefab prefab;

        /// <summary>
        /// The base name of the cluster. I don't think this is actively used at this moment
        /// </summary>
        public string baseName {
            get {
                int colonPositon = this.name.IndexOf(':');
                if (colonPositon < 0)
                    return this.name;
                return this.name[..colonPositon];
            }
        }

        /// <summary>
        /// All cluster instance of a multi-cluster
        /// </summary>
        /// A cluster is a multi-cluster when there is more than one instance.
        /// The actual instances are only created at runtime.
        /// The value instanceCount determines how many instances will be present at runtime.
        //[NonSerialized]
        [SerializeReference]
        public Cluster[] instances;

        /// <summary>
        /// The number of cluster instances in a multi-cluster
        /// </summary>
        /// A cluster is a multi-clsuter when there is more than one instance.
        [SerializeField]
        public int instanceCount = 1;
        /// <summary>
        /// The mapping from things to cluster instances
        /// </summary>
        /// In a multi-cluster each instance can be used for a thing.
        /// Cluster instance may also not (yet) be mapped to a thing.
        public Dictionary<int, Cluster> thingClusters = new();

        /// <summary>
        /// All nuclei in this cluster
        /// </summary>
        [SerializeReference]
        public List<Nucleus> nuclei = new();
        // the nuclei sorted using topological sorting
        // to ensure that the cluster is computer in the right order

        //public List<Nucleus> sortedNuclei;

        #region Init

        /// <summary>
        /// Instantiate a new copy of a ClusterPrefab in the given parent
        /// </summary>
        /// <param name="prefab">The prefab to use</param>
        /// <param name="parent">The cluster in which this new cluster will be placed</param>
        public Cluster(ClusterPrefab prefab, Cluster parent) {
            this.prefab = prefab;
            this.name = prefab.name;

            this.parent = parent;
            this.parent?.nuclei.Add(this);
            ClonePrefab();
            // _ = this.inputs;
            //this.sortedNuclei = TopologicalSort(this.nuclei);
        }

        /// <summary>
        /// Add a new cluster to a ClusterPrefab
        /// </summary>
        /// <param name="prefab">The prefab to copy</param>
        /// <param name="parent">The prefab in which the new copy is placed</param>
        public Cluster(ClusterPrefab prefab, ClusterPrefab parent = null) {
            this.prefab = prefab;
            this.name = prefab.name;
            if (parent != null)
                this.parent = parent.cluster;

            ClonePrefab();
            // _ = this.inputs;
            //this.sortedNuclei = TopologicalSort(this.nuclei);
        }

        /// <summary>
        /// Clone a prefab.
        /// </summary>
        /// Strange that this does not take any parameters or return values.
        /// Where which the clone be found???
        private void ClonePrefab() {
            if (this.prefab == null || this.prefab.cluster == null || this.prefab.cluster.nuclei == null)
                return;
                
            Nucleus[] prefabNuclei = this.prefab.cluster.nuclei.ToArray();

            // first clone the nuclei without their connections
            foreach (Nucleus nucleus in prefabNuclei) {
                nucleus.ShallowCloneTo(this);
            }
            Nucleus[] clonedNuclei = this.nuclei.ToArray();

            // Now clone the connections
            for (int nucleusIx = 0; nucleusIx < prefabNuclei.Length; nucleusIx++) {
                Nucleus prefabNucleus = prefabNuclei[nucleusIx];
                if (prefabNucleus is not Neuron prefabNeuron)
                    continue;

                Nucleus clonedNucleus = clonedNuclei[nucleusIx];
                if (clonedNucleus == null || clonedNucleus is not Neuron clonedNeuron)
                    continue;

                foreach (Synapse prefabSynapse in prefabNeuron.synapses) {
                    Neuron synapseNeuron = prefabSynapse.neuron;
                    if (synapseNeuron.parent.prefab != null && synapseNeuron.parent.prefab != this.prefab) {
                        // Neuron is in another cluster, find the cloned cluster first
                        Cluster prefabCluster = synapseNeuron.parent;
                        Cluster clonedCluster = this.nuclei.Find(n => n.name == prefabCluster.name) as Cluster;
                        if (clonedCluster == null)
                            continue;

                        // Now find the neuron in that cloned cluster
                        int neuronIx = GetNucleusIndex(prefabCluster.nuclei, prefabSynapse.neuron.name);
                        if (neuronIx < 0)
                            // Could not find the neuron in the prefab cluster
                            continue;
                        if (clonedCluster.nuclei[neuronIx] is not Neuron clonedSender)
                            // Could not find the neuron in the cloned cluster
                            continue;

                        clonedSender.AddReceiver(clonedNeuron, prefabSynapse.weight);
                        //Debug.Log($"Add synapse {clonedCluster.name}.{clonedSender.name} -> {clonedNeuron.name} [{clonedSender.receivers.Count}]");
                    }
                    else {
                        int ix = GetNucleusIndex(this.prefab.cluster.nuclei, prefabSynapse.neuron);
                        if (ix < 0)
                            continue;
                        if (clonedNuclei[ix] is not Neuron clonedSender)
                            continue;

                        // Copy the receivers which will also create the synapse
                        clonedSender.AddReceiver(clonedNeuron, prefabSynapse.weight);
                        // Debug.Log($"Add synapse {clonedSender.name} -> {clonedNeuron.name}");
                    }
                }
            }

            //if (Application.isPlaying) {
                // Only create cluster siblings at runtime
                foreach (Nucleus clonedNucleus in clonedNuclei) {
                    if (clonedNucleus is not Cluster clonedCluster)
                        continue;

                    List<Cluster> siblings = new() {
                        clonedCluster
                    };
                    for (int instanceIx = 1; instanceIx < clonedCluster.instanceCount; instanceIx++) {
                        // Create another sibling
                        Debug.Log($"create {clonedCluster.prefab.name} sibling");
                        Cluster sibling = new(clonedCluster.prefab, this) {
                            name = $"{clonedCluster.baseName}: {instanceIx}",
                            parent = this.parent,
                            instanceCount = this.instanceCount,
                        };
                        siblings.Add(sibling);
                        CopyAllExternalReceivers(clonedCluster, sibling, clonedCluster.prefab, this);
                    }
                    Cluster[] siblingClusters = siblings.ToArray();
                    foreach (Cluster sibling in siblings)
                        sibling.instances = siblingClusters;
                }

                // Ensure that all neurons are computed to initialize bias
                foreach (Nucleus clonedNucleus in clonedNuclei) {
                    if (clonedNucleus is not Cluster)
                        clonedNucleus.UpdateStateIsolated();
                }
            //}
        }

        /// \copydoc NanoBrain::Nucleus::ShallowCloneTo
        public override Nucleus ShallowCloneTo(Cluster parent) {
            // Clusters should not be cloned, but instantiated from the prefab....
            Cluster clone = new(this.prefab, parent) {
                name = this.name,
                parent = this.parent,
                instanceCount = this.instanceCount,
            };
            // Somehow siblingClusters should be cloned too. Believe I do this in ClonePrefab right now.

            return clone;
        }

        private static void CopyAllExternalReceivers(Cluster sourceCluster, Cluster sibling, ClusterPrefab prefabParent, Cluster clonedParent) {

            for (int nucleusIx = 0; nucleusIx < sourceCluster.nuclei.Count; nucleusIx++) {
                Nucleus sourceNucleus = sourceCluster.nuclei[nucleusIx];
                if (sourceNucleus is not Neuron sourceNeuron)
                    continue;

                if (sibling.nuclei[nucleusIx] is not Neuron clonedNeuron)
                    continue;

                // copy the receivers (and thus synapses) from the source to the sibling
                foreach (Nucleus receiver in sourceNeuron.receivers) {
                    if (receiver is not Neuron receiverNeuron)
                        continue;

                    int ix = GetNucleusIndex(clonedParent.nuclei, receiver);
                    if (ix < 0 || ix >= clonedParent.nuclei.Count)
                        continue;

                    // Find the synapse for the weight
                    float weight = 1;
                    foreach (Synapse synapse in receiverNeuron.synapses) {
                        // Find the weight for this synapse
                        if (synapse.neuron == sourceNucleus) {
                            weight = synapse.weight;
                            break;
                        }
                    }

                    clonedNeuron.AddReceiver(receiver, weight);
                    Debug.Log($"external: {receiver.name} receives from {clonedNeuron.name} {clonedNeuron.GetHashCode()}");
                }
            }

        }

        /// <summary>
        /// Get the index of a nucleus in a list of nuclei
        /// </summary>
        /// <param name="nuclei">The list of nuclei to search</param>
        /// <param name="nucleus">The nucleus to find</param>
        /// <returns>The index of the nucleus in the list or -1 when it has not been found</returns>
        public static int GetNucleusIndex(List<Nucleus> nuclei, Nucleus nucleus) {
            int i = 0;
            foreach (Nucleus nucleiElement in nuclei) {
                //for (int i = 0; i < nuclei.Length; i++) {
                if (nucleiElement == nucleus)
                    return i;
                i++;
            }
            return -1;
        }

        /// <summary>
        /// Get the index of a nucleus with the given name in a list of nuclei
        /// </summary>
        /// <param name="nuclei">The list of nuclei to search</param>
        /// <param name="nucleusName">The name of the nucleus to find</param>
        /// <returns>The index of the nucleus in the list or -1 when it has not been found</returns>
        public static int GetNucleusIndex(List<Nucleus> nuclei, string nucleusName) {
            int i = 0;
            foreach (Nucleus nucleiElement in nuclei) {
                //for (int i = 0; i < nuclei.Length; i++) {
                if (nucleiElement.name == nucleusName)
                    return i;
                i++;
            }
            return -1;
        }

        #endregion Init

        #region Cluster Array

        /// <summary>
        /// Increase the number of instances in an multi-cluster
        /// </summary>
        /// /remark Note this does not create the instances.
        /// This is only intended to be used for prefabs.
        public void AddInstance() {
            this.instanceCount++;
        }

        /// <summary>
        /// Create an new instance in a multi-cluster
        /// </summary>
        /// <param name="prefab">The prefab to use to create the new instance</param>
        /// /remark This does not change the instanceCount.
        /// It should only be used at runtime.
        public void AddInstance(ClusterPrefab prefab) {
            // Ensure siblingClusters exists
            if (this.instances == null || this.instances.Length == 0)
                this.instances = new Cluster[1] { this };

            // Prepare the new array
            int newLength = this.instances.Length + 1;
            Cluster[] newSiblings = new Cluster[newLength];

            for (int i = 0; i < newSiblings.Length - 1; i++)
                newSiblings[i] = this.instances[i];

            //Cluster newCluster = this.Clone(prefab) as Cluster;
            Cluster newCluster = new(prefab);
            string baseName = this.name;
            int colonPos = baseName.IndexOf(":");
            if (colonPos > 0)
                baseName = baseName[..colonPos];
            newCluster.name = $"{baseName}: {newLength - 1}";
            newSiblings[newLength - 1] = newCluster;

            // All siblingClusters need to user this array!
            foreach (Cluster sibling in newSiblings)
                sibling.instances = newSiblings;
        }

        /// <summary>
        /// Decrease the number of instance in a multi-cluster
        /// </summary>
        public void RemoveInstance() {
            if (instanceCount > 1)
                instanceCount--;
            else {
                // It is not clear to me why we update the siblingClusters when the 
                // instanceCount <= 1....
                if (this.instances == null || this.instances.Length <= 1)
                    return;

                // Prepare the new array
                int newLength = this.instances.Length - 1;
                Cluster[] newClusters = new Cluster[newLength];

                for (int i = 0; i < newLength; i++)
                    newClusters[i] = this.instances[i];

                Neuron.Delete(this.instances[^1]);
                this.instances = newClusters;
            }
        }

        /// <summary>
        /// Remove a mapping from a thing to a cluster such that it becomes available for new things
        /// </summary>
        /// <param name="cluster">The multi-cluster instance which not no longer be mapped</param>
        private void RemoveThingCluster(Cluster cluster) {
            List<int> keysToRemove = new();
            foreach (KeyValuePair<int, Cluster> kvp in thingClusters) {
                if (kvp.Value == cluster)
                    keysToRemove.Add(kvp.Key);
            }

            foreach (int thingId in keysToRemove)
                thingClusters.Remove(thingId);
        }

        #endregion ClusterArray

        /// <summary>
        /// This gives the order in which nuclei should be computed when a nucleus is updated
        /// </summary>
        [NonSerialized]
        private Dictionary<Nucleus, List<Nucleus>> _computeOrders;
        /// <summary>
        /// This gives the order in which nuclei should be computed when a nucleus is updated
        /// </summary>
        public Dictionary<Nucleus, List<Nucleus>> computeOrders {
            get {
                if (_computeOrders == null || _computeOrders.Count == 0) {
                    _computeOrders = new();
                    foreach (Nucleus nucleus in this.nuclei)
                        _computeOrders[nucleus] = TopologicalSort2(nucleus);
                }
                return _computeOrders;
            }
        }
        /// <summary>
        /// Refresh the order in which neurons should be computed
        /// </summary>
        public void RefreshComputeOrders() {
            this._computeOrders = null;
        }

        private List<Nucleus> TopologicalSort2(Nucleus startNode) {
            Dictionary<Nucleus, int> inDegree = new();
            //HashSet<Nucleus> visited = new();

            // Calculate in-degrees for all nodes reachable from the start node
            Queue<Nucleus> queue = new();
            queue.Enqueue(startNode);
            //visited.Add(startNode);
            inDegree[startNode] = 0;

            while (queue.Count > 0) {
                Nucleus current = queue.Dequeue();
                List<Nucleus> receivers = null;
                if (current is Neuron neuron)
                    receivers = neuron.receivers;
                else if (current is Cluster cluster)
                    receivers = cluster.CollectReceivers();

                foreach (Nucleus receiver in receivers) {
                    if (!inDegree.ContainsKey(receiver)) {
                        //visited.Add(receiver);
                        inDegree[receiver] = 0;
                        queue.Enqueue(receiver);
                    }
                    inDegree[receiver]++;
                }
            }

            // Perform topological sort on all reachable nodes
            queue.Clear();
            foreach (Nucleus node in inDegree.Keys) {
                if (inDegree[node] == 0)
                    queue.Enqueue(node);
            }

            List<Nucleus> sortedOrder = new();
            while (queue.Count > 0) {
                Nucleus current = queue.Dequeue();
                sortedOrder.Add(current); // Process the node

                List<Nucleus> receivers = null;
                if (current is Neuron neuron)
                    receivers = neuron.receivers;
                else if (current is Cluster cluster)
                    receivers = cluster.CollectReceivers();

                foreach (Nucleus receiver in receivers) {
                    if (inDegree.ContainsKey(receiver)) {
                        inDegree[receiver]--;
                        if (inDegree[receiver] == 0) // If all dependencies resolved
                            queue.Enqueue(receiver);
                    }
                }
            }

            // Check for cycles in the graph
            // if (sortedOrder.Count != visited.Count)
            //     throw new InvalidOperationException("Graph is not a DAG; a cycle exists.");

            return sortedOrder;
        }

        /// <summary>
        /// The first nucleus in a cluster is the default output
        /// </summary>
        public virtual Neuron defaultOutput {//=> this.nuclei[0] as Nucleus;
            get {
                if (this.nuclei.Count > 0)
                    return this.nuclei[0] as Neuron;
                return null;
            }
        }

        /// <summary>
        /// The neurons without outgoing connections
        /// </summary>
        /// These neurons can potentially be connected to neurons in other clusters
        [NonSerialized]
        protected List<Neuron> _outputs = null;
        /// <summary>
        /// The neurons without outgoing connections
        /// </summary>
        /// These neurons can potentially be connected to neurons in other clusters
        public List<Neuron> outputs {
            get {
                if (this._outputs == null || this._outputs.Count == 0) {
                    this._outputs = new();
                    foreach (Nucleus nucleus in this.nuclei) {
                        if (nucleus is Neuron neuron && neuron.receivers.Count == 0)
                            this._outputs.Add(neuron);
                    }
                }
                return this._outputs;
            }
        }
        /// <summary>
        /// Reset the list of outputs such that they will be re-determined
        /// </summary>
        public void RefreshOutputs() {
            this._outputs = null;
        }

        /// <summary>
        /// Try to find a nucleus in this cluster
        /// </summary>
        /// <param name="nucleusName">The name of the nucleus to find</param>
        /// <param name="foundNucleus">The found nucleus or null if it is not found</param>
        /// <returns>True when the nucleus is found, false otherwise</returns>
        public bool TryGetNucleus(string nucleusName, out Nucleus foundNucleus) {
            foreach (Nucleus receptor in this.nuclei) {
                if (receptor is Nucleus nucleus)
                    if (nucleus.name == nucleusName) {
                        // if (nucleus is Cluster cluster)
                        //     cluster.CheckInstances();
                        foundNucleus = nucleus;
                        return true;
                    }
            }
            foundNucleus = null;
            return false;
        }

        /// <summary>
        /// Get a nucleus in this cluster
        /// </summary>
        /// <param name="nucleusName">The name of the nucleus to find</param>
        /// <returns>The found nucleus or null when it is not found</returns>
        public Nucleus GetNucleus(string nucleusName) {
            int dotPosition = nucleusName.IndexOf('.');
            if (dotPosition >= 0) {
                string clusterName = nucleusName[..dotPosition];
                string clusterName0 = clusterName + ": 0";
                foreach (Nucleus nucleus in this.nuclei) {
                    if (nucleus is Cluster cluster) {
                        if (cluster.name == clusterName || cluster.name == clusterName0) {
                            // cluster.CheckInstances();
                            string subNucleusName = nucleusName[(dotPosition + 1)..];
                            return cluster.GetNucleus(subNucleusName);
                        }
                    }
                }
                return null;
            }
            else {
                string nucleusName0 = nucleusName + ": 0";
                foreach (Nucleus nucleus in this.nuclei) {
                    if (nucleus is Cluster cluster) {
                        if (nucleus.name == nucleusName || nucleus.name == nucleusName0) {
                            // cluster.CheckInstances();
                            return nucleus;
                        }
                    }
                    else if (nucleus.name == nucleusName)
                        return nucleus;
                }
                return null;
            }
        }

        /// <summary>
        /// Get a neuron in this cluster
        /// </summary>
        /// <param name="neuronName">The name of the neuron to find</param>
        /// <returns>The found neuron or null when it is not found</returns>
        public Neuron GetNeuron(string neuronName) {
            if (this.nuclei == null)
                return null;

            foreach (Nucleus nucleus in this.nuclei) {
                if (nucleus is Neuron neuron && neuron.name == neuronName)
                    return neuron;
            }
            return null;
        }

        /// <summary>
        /// Get a neuron in an instance of a multi-cluster
        /// </summary>
        /// <param name="thingId">The id of the thing mapped to the cluster instance</param>
        /// <param name="neuronName">The name of the neuron to find</param>
        /// <param name="thingName">The name of the thing mapped to the cluster instance</param>
        /// <returns>The found neuron or null when it is not found</returns>
        /// The cluster instance mapped to the thing will be neuron.parent if a neuron is found.
        public Neuron GetNeuron(int thingId, string neuronName, string thingName = null) {
            if (this.instances == null || this.instances.Length <= 1)
                return this.GetNeuron(neuronName);

            // See if we are already using a cluster for thingId
            thingClusters ??= new();
            if (thingClusters.TryGetValue(thingId, out Cluster cluster))
                return cluster.GetNeuron(neuronName);

            // Find the cluster with the lowest value neuron
            Neuron lowestNeuron = null;
            foreach (Cluster sibling in this.instances) {
                Neuron neuron = sibling.GetNeuron(neuronName);
                if (lowestNeuron == null || neuron.outputMagnitude < lowestNeuron.outputMagnitude)
                    lowestNeuron = neuron;
            }
            Cluster selectedCluster = lowestNeuron.parent;
            RemoveThingCluster(selectedCluster);
            selectedCluster.name = baseName + ": " + thingName;
            thingClusters[thingId] = selectedCluster;
            return lowestNeuron;
        }

        /// <summary>
        /// Delete a nucleus from this clsuter
        /// </summary>
        /// <param name="nucleus">The nucleus to delete</param>
        /// <returns>True if a nucleus was deleted, false if the nucleus could not be found</returns>
        public bool DeleteNucleus(Nucleus nucleus) {
            if (this.nuclei.Contains(nucleus) == false) {
                // Try to find the nucleus by name
                if (TryGetNucleus(nucleus.name, out nucleus) == false)
                    return false;
            }

            Neuron.Delete(nucleus);
            //int nucleusIx = this.nuclei.IndexOf(nucleus);
            this.nuclei.Remove(nucleus);
            //this.prefab.cluster.nuclei.RemoveAt(nucleusIx);
            RefreshOutputs();

            return true;
        }

        #region Receivers

        /// <summary>
        /// Collect all receiving nuclei of signals from this cluster
        /// </summary>
        /// <param name="removeDuplicates">Ensure that a receiver is only listed once in the result</param>
        /// <returns>The list of receivers</returns>
        public virtual List<Nucleus> CollectReceivers(bool removeDuplicates = false) {
            List<Nucleus> receivers = new();
            foreach (Nucleus outputNucleus in this.nuclei) {
                if (outputNucleus is not Neuron output)
                    continue;

                // Debug.Log($"output {this.name} {outputNucleus.name}");
                foreach (Nucleus receiver in output.receivers) {
                    // Debug.Log($"output    {receiver.name}");
                    // Only add receivers outside this cluster                    
                    if (receiver.parent.prefab != this.prefab) {
                        if (removeDuplicates == false || receivers.Contains(receiver) == false)
                            // Debug.Log($"   YES");
                            receivers.Add(receiver);
                    }
                }
            }
            return receivers;
        }

        /// <summary>
        /// Collect all synapses of senders in another cluster of signals to this cluster
        /// </summary>
        /// <param name="otherCluster">The other cluster with sending neurons</param>
        /// <returns>A list of synapses to the neurons in the other clusters</returns>
        public List<Synapse> CollectSynapsesTo(Cluster otherCluster) {
            List<Synapse> collectedSynapses = new();

            foreach (Nucleus nucleus in this.nuclei) {
                if (nucleus is not Neuron neuron)
                    continue;
                foreach (Synapse synapse in neuron.synapses) {
                    if (synapse.neuron.parent == otherCluster)
                        collectedSynapses.Add(synapse);
                }
            }
            return collectedSynapses;
        }

        #endregion Receivers

        #region Update

        /// <summary>
        /// Update the state of the nucleus and all nuclei receiving from it
        /// </summary>
        /// <param name="startNucleus">The nucleus to start updating</param>
        public void UpdateFromNucleus(Nucleus startNucleus) {
            // no bias+synapse input state calculation for now...

            if (this.computeOrders.ContainsKey(startNucleus) == false) {
                Debug.LogError($"{this.name} compute orders does not contain an order for {startNucleus.name}");
                return;
            }

            List<Nucleus> computeOrder = this.computeOrders[startNucleus];
            foreach (Nucleus nucleus in computeOrder) {
                if (nucleus is not Cluster) {
                    nucleus.UpdateStateIsolated();
                    if (nucleus is Neuron neuron) {
                        foreach (Nucleus receiver in neuron.receivers) {
                            if (receiver.parent != this) {
                                //Debug.Log($"   External: {receiver.parent.name}.{receiver.name}");
                                receiver.parent.UpdateFromNucleus(receiver);
                            }
                        }
                    }
                }
            }
        }

        /// \copydoc NanoBrain::Nucleus::UpdateStateIsolated
        public override void UpdateStateIsolated() {
            throw new Exception("Cluster should not be updated!");
        }

        #endregion Update

        /// <summary>
        /// Recalculate derived properties
        /// </summary>
        /// This can be used to recalculate derived properties after the set of nuclei has been changed
        public void Refresh() {
            // This should not be needed, but somehow somewhere the parent is changed...
            foreach (Nucleus nucleus in this.nuclei) {
                nucleus.parent = this;
            }
            RefreshOutputs();
            RefreshComputeOrders();
        }

    }

}