using System.Collections.Generic;
using System.Linq;
using UnityEditor;

using UnityEngine;
using UnityEngine.UIElements;

namespace NanoBrain {

    [CustomEditor(typeof(ClusterPrefab))]
    public class ClusterEditor : ClusterViewer {

        public override VisualElement CreateInspectorGUI() {
            ClusterPrefab prefab = target as ClusterPrefab;
            if (prefab != null)
                prefab.EnsureInitialization();

            serializedObject.Update();

            VisualElement root = new();
            CreateEditor(root, prefab, null);

            serializedObject.ApplyModifiedProperties();
            return root;
        }

        public GraphView CreateEditor(VisualElement root, ClusterPrefab cluster, GameObject gameObject) {
            root.style.paddingLeft = 0;
            root.style.paddingRight = 0;
            root.style.paddingTop = 0;
            root.style.paddingBottom = 0;

            root.styleSheets.Add(Resources.Load<StyleSheet>("GraphStyles"));

            VisualElement mainContainer = new() {
                style = {
                flexDirection = FlexDirection.Row,
                }
            };
            GraphEditor graphContainer = new(cluster);
            graphContainer.style.flexShrink = 0;
            graphContainer.style.width = 300;
            graphContainer.style.overflow = Overflow.Hidden;

            VisualElement inspectorContainer = new() {
                name = "inspector",
                style = {
                    minHeight = 450,
                    width = 300,
                    flexGrow = 0,
                    flexDirection = FlexDirection.Row,
                }
            };

            mainContainer.Add(graphContainer);
            mainContainer.Add(inspectorContainer);
            root.Add(mainContainer);

            graphContainer.SetGraph(gameObject, inspectorContainer);

            return graphContainer;
        }

        public class GraphEditor : GraphView {

            protected ClusterPrefab prefab;

            public GraphEditor(ClusterPrefab prefab) : base(prefab.output.parent) {
                this.prefab = prefab;

                // In a Prefab editor, no instance exists but we need it for the ClusterViewer.
                // So we create a temporary instance
                Cluster cluster = new(prefab);
                this.currentCluster = cluster;

                Button addButton = new(() => OnAddClusterOutput()) {
                    text = "Add"
                };
                topMenuContainer?.Add(addButton);

                Add(topMenuContainer);
            }

            void OnAddClusterOutput() {
                Nucleus newOutput = new Neuron(this.prefab, "New Output");
                this.prefab.RefreshOutputs();
                outputsPopup.choices = this.prefab.outputs.Select(output => output.name).ToList();
                outputsPopup.value = newOutput.name;

                this.currentNucleus = newOutput;
            }

            public void SetGraph(GameObject gameObject, VisualElement inspectorContainer) {
                this.gameObject = gameObject;

                if (Application.isPlaying == false)
                    this.serializedBrain = new SerializedObject(this.prefab);
                this.selectedOutput = this.currentCluster.outputs[0];
                this.currentNucleus = this.selectedOutput;
                //this.currentCluster = this.currentNucleus.parent;
                Rebuild(inspectorContainer);
                // if (outputsPopup != null)
                //     OnOutputChanged(outputsPopup.choices[0]);
            }

            private void Rebuild(VisualElement inspectorContainer) {

                if (this.currentNucleus == null) {
                    inspectorContainer.Clear();
                    return;
                }

                string path = AssetDatabase.GetAssetPath(this.prefab); // or known path
                this.prefabAsset = AssetDatabase.LoadAssetAtPath<ClusterPrefab>(path);
                if (this.prefabAsset == null) {
                    // create in memory save if it doesn't exist
                    this.prefabAsset = CreateInstance<ClusterPrefab>();
                    //Debug.LogError("Cluster Prefab is not found on disk");
                }
                // DrawInspector(inspectorContainer);
                if (inspectorContainer == null)
                    return;

                inspectorContainer.Clear();
                if (this.currentNucleus == null)
                    return;

                // create a SerializedObject wrapper so Unity inspector controls work (and Undo)
                SerializedObject so = new(prefabAsset);

                foreach (Nucleus nucleus in this.prefab.nuclei) {
                    nucleus.Initialize();
                }

                this.inspectorIMGUIContainer = new IMGUIContainer(() => InspectorHandler(so));

                inspectorContainer.Add(inspectorIMGUIContainer);
            }

            #region Inspector

            private VisualElement inspectorIMGUIContainer;
            private bool showSynapses = true;
            private bool showActivation = true;
            protected bool breakOnWake = false;
            protected bool trace = false;

            void InspectorHandler(SerializedObject serializedObject) {
                bool anythingChanged = false;

                if (serializedObject == null || serializedObject.targetObject == null)
                    return;

                if (this.currentNucleus == null)
                    return;

                serializedObject.Update();

                GUIStyle headerStyle = new(EditorStyles.boldLabel) {
                    alignment = TextAnchor.MiddleLeft,
                    margin = new RectOffset(10, 0, 4, 4)
                };
                GUIStyle boldTextFieldStyle = new(EditorStyles.textField) {
                    fontStyle = FontStyle.Bold
                };

                // Nucleus type
                string nucleusType = this.currentNucleus.GetType().Name;
                GUILayout.Label(nucleusType, headerStyle);

                // Nucleus name
                if (this.currentNucleus.parent is Cluster parentCluster) {
                    EditorGUILayout.BeginHorizontal();
                    if (GUILayout.Button(this.currentNucleus.parent.name))
                        OnClusterClick(parentCluster);
                    EditorGUI.BeginDisabledGroup(true);
                    EditorGUILayout.TextField(this.currentNucleus.name, boldTextFieldStyle);
                    EditorGUI.EndDisabledGroup();
                    if (GUILayout.Button("Reimport"))
                        ReimportCluster(parentCluster);
                    EditorGUILayout.EndHorizontal();
                }
                else {
                    string newName = EditorGUILayout.TextField(this.currentNucleus.name, boldTextFieldStyle);
                    if (newName != this.currentNucleus.name) {
                        this.currentNucleus.name = newName;
                        this.prefab.RefreshOutputs();
                        outputsPopup.choices = this.prefab.outputs.Select(output => output.name).ToList();
                        anythingChanged = true;
                    }
                }

                // Current output value
                if (Application.isPlaying) {
                    if (currentNucleus is Neuron currentNeuron1) {
                        GUIContent nameLabel = new("Output", currentNeuron1.outputValue.ToString());
                        EditorGUILayout.FloatField(nameLabel, currentNeuron1.outputMagnitude);
                    }
                    else
                        EditorGUILayout.LabelField(" ");
                }
                else
                    EditorGUILayout.LabelField(" ");

                // Memory cell
                if (this.currentNucleus is MemoryCell memory)
                    MemoryCellInspector(memory, ref anythingChanged);
                // Cluster
                else if (this.currentNucleus is Cluster cluster)
                    ClusterInspector(cluster, ref anythingChanged);
                // Other
                else
                    NucleusInspector(this.currentNucleus, ref anythingChanged);

                if (GUILayout.Button("Delete"))
                    DeleteNucleus(this.currentNucleus);

                serializedObject.ApplyModifiedProperties();
                if (anythingChanged) {
                    EditorUtility.SetDirty(prefabAsset);
                    AssetDatabase.SaveAssets();
                }
            }

            protected void MemoryCellInspector(MemoryCell memoryCell, ref bool anythingChanged) {
                memoryCell.staticMemory = EditorGUILayout.Toggle("Static Memory", memoryCell.staticMemory);
                NucleusInspector(memoryCell, ref anythingChanged);
            }

            protected void ClusterInspector(Cluster cluster, ref bool anythingChanged) {
                EditorGUILayout.BeginHorizontal();

                int instanceCount = cluster.instanceCount;
                if (instanceCount <= 1) {
                    if (cluster.siblingClusters != null && cluster.siblingClusters.Length > 1)
                        instanceCount = cluster.siblingClusters.Count();
                    else
                        instanceCount = 1;
                }
                EditorGUILayout.IntField("Instances", instanceCount, GUILayout.MinWidth(150));

                if (GUILayout.Button("Add")) {
                    Undo.RecordObject(prefabAsset, "Array add " + prefabAsset.name);
                    //cluster.AddInstance(this.prefab);
                    cluster.AddInstance();
                    anythingChanged = true;
                }
                if (GUILayout.Button("Del")) {
                    Undo.RecordObject(prefabAsset, "Array delete " + prefabAsset.name);
                    cluster.RemoveInstance();
                    anythingChanged = true;
                }
                EditorGUILayout.EndHorizontal();

                if (GUILayout.Button("Reimport Cluster"))
                    ReimportCluster(cluster);
            }

            protected void NucleusInspector(Nucleus nucleus, ref bool anythingChanged) {
                SynapsesInspector(ref anythingChanged);
                ActivationInspector(ref anythingChanged);

                EditorGUILayout.Space();
                breakOnWake = EditorGUILayout.Toggle("Break on wake", breakOnWake);
                if (breakOnWake && this.currentNucleus is Neuron currentNeuron) {
                    if (currentNeuron.isSleeping == false)
                        Debug.Break();
                }
                trace = EditorGUILayout.Toggle("Trace", trace);
                this.currentNucleus.trace = trace;
            }

            protected void SynapsesInspector(ref bool anythingChanged) {
                showSynapses = EditorGUILayout.BeginFoldoutHeaderGroup(showSynapses, "Synapses");
                if (showSynapses) {
                    if (this.currentNucleus is Neuron neuron2) {
                        Neuron.CombinatorType newCombinator = (Neuron.CombinatorType)EditorGUILayout.EnumPopup("Combinator", neuron2.combinator);
                        anythingChanged |= newCombinator != neuron2.combinator;
                        neuron2.combinator = newCombinator;
                    }

                    EditorGUIUtility.wideMode = true;
                    float previousLabelWidth = EditorGUIUtility.labelWidth;
                    EditorGUIUtility.labelWidth = 100;

                    Vector3 newBias = EditorGUILayout.Vector3Field("Bias", this.currentNucleus.bias);
                    anythingChanged |= newBias != this.currentNucleus.bias;
                    this.currentNucleus.bias = newBias;
                    EditorGUIUtility.labelWidth = previousLabelWidth;

                    Nucleus[] array = null;
                    int elementIx = -1;
                    if (this.currentNucleus.synapses.Count > 0) {
                        Synapse[] synapses = this.currentNucleus.synapses.ToArray();
                        foreach (Synapse synapse in synapses) {
                            if (synapse.neuron == null)
                                continue;

                            if (array != null) {
                                if (synapse.neuron.parent is Cluster iCluster && elementIx > 0) {
                                    int thisElementIx = Cluster.GetNucleusIndex(iCluster.clusterNuclei, synapse.neuron);
                                    if (thisElementIx == elementIx)
                                        continue;
                                    else
                                        elementIx = thisElementIx;
                                }
                                if (array.Contains(synapse.neuron))
                                    continue;
                                else if (array.Contains(synapse.neuron.parent))
                                    continue;
                            }
                            else {
                                if (synapse.neuron.parent is Cluster iReceptor) {
                                    array = iReceptor.siblingClusters;
                                    if (iReceptor is Cluster iCluster)
                                        elementIx = Cluster.GetNucleusIndex(iCluster.clusterNuclei, synapse.neuron);
                                }
                            }

                            EditorGUILayout.Space();

                            if (Application.isPlaying) {
                                if (synapse.neuron is Neuron synapseNeuron) {
                                    Vector3 value = synapseNeuron.outputValue * synapse.weight;
                                    GUIContent synapseValueLabel = new(synapse.neuron.name, synapseNeuron.outputValue.ToString());
                                    EditorGUILayout.FloatField(synapseValueLabel, synapseNeuron.outputMagnitude);
                                }
                            }
                            else {
                                EditorGUILayout.BeginHorizontal();

                                if (synapse.neuron.clusterPrefab != this.currentNucleus.clusterPrefab) {
                                    // If it is a cluster
                                    GUIStyle labelStyle = new(GUI.skin.label);
                                    float labelWidth = 200;
                                    if (synapse.neuron.clusterPrefab != null) {
                                        labelWidth = labelStyle.CalcSize(new GUIContent($"{synapse.neuron.clusterPrefab.name}.")).x;
                                        GUILayout.Label($"{synapse.neuron.clusterPrefab.name}", GUILayout.Width(labelWidth));
                                    }
                                    //string[] options = synapse.neuron.parent.clusterNuclei.Select(n => n.name).ToArray();
                                    string[] options = synapse.neuron.clusterPrefab.nuclei.Select(n => n.name).ToArray();
                                    int selectedIndex = System.Array.IndexOf(options, synapse.neuron.name);
                                    int newIndex = EditorGUILayout.Popup(selectedIndex, options);
                                    // if (newIndex != selectedIndex && synapse.neuron.clusterPrefab.nuclei[newIndex] is Neuron newNeuron)
                                    //     ChangeSynapse(synapse, newNeuron);
                                    if (newIndex != selectedIndex) {
                                        // It shall be ensured that the parent.clusterNuclei and
                                        // clusterPrefab.nuclei contain the same neurons in the same order....
                                        Nucleus selectedNucleus = synapse.neuron.parent.clusterNuclei[newIndex];
                                        Neuron newNeuron = selectedNucleus as Neuron;
                                        ChangeSynapse(synapse, newNeuron);
                                    }
                                }
                                else
                                    GUILayout.Label(synapse.neuron.name);

                                bool disconnecting = GUILayout.Button("Disconnect", GUILayout.Width(80));
                                if (disconnecting && synapse.neuron is Neuron synapseNeuron) {
                                    synapseNeuron.RemoveReceiver(this.currentNucleus);
                                    this.prefab.GarbageCollection();
                                    anythingChanged = true;
                                }
                                EditorGUILayout.EndHorizontal();

                            }

                            EditorGUI.indentLevel++;
                            float newWeight = EditorGUILayout.FloatField("Weight", synapse.weight);
                            if (newWeight != synapse.weight) {
                                // if (synapse.neuron.parent is IReceptor receptor) {
                                //     Nucleus[] receptorArray = receptor.nucleiArray;
                                //     foreach (Synapse s in this.currentNucleus.synapses) {
                                //         if (s.neuron.parent is IReceptor r && r.nucleiArray == receptorArray)
                                //             s.weight = newWeight;
                                //     }
                                // }
                                // else
                                synapse.weight = newWeight;
                                anythingChanged = true;
                            }
                            EditorGUI.indentLevel--;
                        }
                    }

                    EditorGUILayout.Space();
                    anythingChanged |= ConnectNucleus(this.prefab, this.currentNucleus);
                    anythingChanged |= AddSynapse(this.prefab, this.currentNucleus);
                }
                EditorGUILayout.EndFoldoutHeaderGroup();
            }

            protected void ActivationInspector(ref bool anythingChanged) {
                EditorGUILayout.Space();
                showActivation = EditorGUILayout.BeginFoldoutHeaderGroup(showActivation, "Activation");
                if (showActivation) {
                    if (this.currentNucleus is Neuron neuron) {
                        if (this.currentNucleus is not MemoryCell) {
                            EditorGUILayout.BeginHorizontal();
                            EditorGUILayout.LabelField("Activation Curve", GUILayout.MinWidth(60));
                            if (neuron.curveMax > 0)
                                EditorGUILayout.CurveField(neuron.curve, Color.cyan, new Rect(0, 0, 1, neuron.curveMax), GUILayout.Width(40));
                            else
                                EditorGUILayout.CurveField(neuron.curve, Color.cyan, new Rect(0, neuron.curveMax, 1, -neuron.curveMax), GUILayout.Width(40));
                            Neuron.ActivationType newPreset = (Neuron.ActivationType)EditorGUILayout.EnumPopup(neuron.curvePreset, GUILayout.MinWidth(50));
                            anythingChanged |= newPreset != neuron.curvePreset;
                            neuron.curvePreset = newPreset;
                            EditorGUILayout.EndHorizontal();
                        }
                        // if (neuron is Receptor receptor2) {
                        //     if (receptor2.nucleiArray == null || receptor2.nucleiArray.Count() == 0)
                        //         receptor2.array = new NucleusArray(neuron);
                        // }
                    }

                    EditorGUILayout.Space();
                }
                EditorGUILayout.EndFoldoutHeaderGroup();
            }

            #region Synapses

            protected virtual void AddInput(Nucleus.Type selectedType, Nucleus nucleus) {
                switch (selectedType) {
                    case Nucleus.Type.Neuron:
                        AddNeuronInput(nucleus);
                        break;
                    case Nucleus.Type.MemoryCell:
                        AddMemoryCellInput(nucleus);
                        break;
                    case Nucleus.Type.Cluster:
                        AddClusterInput(nucleus);
                        break;
                    // case Nucleus.Type.Receptor:
                    //     AddReceptorInput(nucleus);
                    //     break;
                    // case Nucleus.Type.ClusterReceptor:
                    //     AddClusterReceptorInput(nucleus);
                    //     break;
                    // case Nucleus.Type.ClusterArray:
                    //     AddClusterArrayInput(nucleus);
                    //     break;
                    default:
                        break;
                }
            }

            protected virtual void AddNeuronInput(Nucleus nucleus) {
                Neuron newNeuroid = new(this.prefab, "New neuron");
                newNeuroid.AddReceiver(nucleus);
                this.currentNucleus = newNeuroid;
            }

            protected virtual void AddMemoryCellInput(Nucleus nucleus) {
                MemoryCell newMemory = new(this.prefab, "New memory cell");
                newMemory.AddReceiver(nucleus);
                this.currentNucleus = newMemory;
            }

            protected virtual void AddClusterInput(Nucleus nucleus) {
                ClusterPickerWindow.ShowPicker(brain => OnClusterPicked(nucleus, brain), "Select Cluster");
            }
            private void OnClusterPicked(Nucleus nucleus, ClusterPrefab selectedPrefab) {
                Cluster subclusterInstance = new(selectedPrefab, this.prefab);
                subclusterInstance.defaultOutput.AddReceiver(nucleus);
            }

            private void ReimportCluster(Cluster subCluster) {
                if (subCluster.siblingClusters == null || subCluster.siblingClusters.Length <= 0) {
                    Cluster reimportedCluster = new(subCluster.prefab, this.prefab);
                    subCluster.MoveReceivers(reimportedCluster);
                    // subcluster should be garbage now...
                    this.currentNucleus = reimportedCluster;
                }
                else {
                    this.currentNucleus = null;
                    List<Cluster> newSiblingsList = new();
                    foreach (Cluster sibling in subCluster.siblingClusters) {
                        Cluster reimportedCluster = new(sibling.prefab, this.prefab) {
                            name = sibling.name
                        };
                        sibling.MoveReceivers(reimportedCluster);
                        newSiblingsList.Add(reimportedCluster);
                        // make the first reimportedCluster the new current nucleus
                        this.currentNucleus ??= reimportedCluster;
                    }
                    Cluster[] newSiblings = newSiblingsList.ToArray();
                    foreach (Cluster sibling in newSiblings)
                        sibling.siblingClusters = newSiblings;
                }
            }

            int selectedConnectNucleus = -1;
            // Connect to another nucleus
            protected virtual bool ConnectNucleus(ClusterPrefab cluster, Nucleus nucleusToConnect) {
                if (cluster == null)
                    return false;

                IEnumerable<Nucleus> synapseNuclei = this.currentNucleus.synapses
                    .Where(synapse => synapse.neuron != null)
                    .Select(synapse => synapse.neuron);

                IEnumerable<Nucleus> nuclei = cluster.nuclei
                    .Except(synapseNuclei);
                IEnumerable<string> nucleiNames = nuclei
                    .Select(n => {
                        int idx = n.name.IndexOf(':');
                        return idx < 0 ? n.name : n.name[..idx];
                    })
                    .Distinct();

                string[] names = nucleiNames.ToArray();
                EditorGUILayout.BeginHorizontal();
                selectedConnectNucleus = EditorGUILayout.Popup(selectedConnectNucleus, names);
                bool connecting = GUILayout.Button("Connect", GUILayout.Width(80));
                EditorGUILayout.EndHorizontal();
                if (connecting) {
                    Nucleus nucleus = nuclei.ElementAt(selectedConnectNucleus);
                    if (nucleus is Cluster subCluster)
                        subCluster.AddArrayReceiver(this.currentNucleus);
                    else if (nucleus is Neuron neuron)
                        neuron.AddReceiver(this.currentNucleus);
                }
                return connecting;
            }

            protected virtual void DeleteNucleus(Nucleus nucleus) {
                if (nucleus == null)
                    return;

                if (nucleus is Neuron neuron) {
                    foreach (Nucleus receiver in neuron.receivers) {
                        if (receiver != null) {
                            this.currentNucleus = receiver;
                            break;
                        }
                    }
                }
                this.prefab.nuclei.Remove(nucleus);

                if (outputsPopup.value == nucleus.name) {
                    this.prefab.RefreshOutputs();
                    outputsPopup.choices = this.prefab.outputs.Select(output => output.name).ToList();
                    outputsPopup.index = 0;
                }

                Neuron.Delete(nucleus);

                this.currentNucleus = this.prefab.output;
            }

            Nucleus.Type selectedType = Nucleus.Type.None;
            protected virtual bool AddSynapse(ClusterPrefab cluster, Nucleus nucleus) {
                if (cluster == null)
                    return false;

                EditorGUILayout.BeginHorizontal();
                selectedType = (Nucleus.Type)EditorGUILayout.EnumPopup(selectedType);
                bool connecting = GUILayout.Button("Add", GUILayout.Width(80));
                EditorGUILayout.EndHorizontal();

                if (connecting) {
                    AddInput(selectedType, this.currentNucleus);
                }
                return connecting;
                // if (selectedType == Nucleus.Type.None)
                //     return false;

                // AddInput(selectedType, this.currentNucleus);
                // return true;
            }

            protected virtual void ChangeSynapse(Synapse synapse, Neuron newNucleus) {
                Neuron synapseNeuron = synapse.neuron as Neuron;
                if (synapse.neuron.parent is Cluster subCluster && subCluster.prefab != this.prefab) {
                    // if (synapse.neuron.parent is ClusterReceptor receptor) {
                    //     // the new nucleus is part of a (cluster) receptor,
                    //     // so we have to change all synapses to this nucleus array elements
                    //     int oldNucleusIx = Cluster.GetNucleusIndex(subCluster.clusterNuclei, synapse.neuron);
                    //     int newNucleusIx = Cluster.GetNucleusIndex(subCluster.clusterNuclei, newNucleus);
                    //     foreach (Nucleus element in receptor.nucleiArray) {
                    //         if (element is not ClusterReceptor clusterReceptor)
                    //             continue;
                    //         // Get the same neuron as the synapse.nucleus in a different element
                    //         // of the ClusterReceptor array
                    //         Nucleus oldElementNucleus = clusterReceptor.clusterNuclei[oldNucleusIx];
                    //         if (oldElementNucleus is not Neuron oldElementNeuron)
                    //             continue;
                    //         // Get the same neuron as newNucleus in a different element
                    //         // of the ClusterReceptor array
                    //         Nucleus newElementNucleus = clusterReceptor.clusterNuclei[newNucleusIx];
                    //         if (newElementNucleus is not Neuron newElementNeuron)
                    //             continue;

                    //         oldElementNeuron.RemoveReceiver(this.currentNucleus);
                    //         newElementNeuron.AddReceiver(this.currentNucleus);
                    //         // Now find the synapse which pointed to the old Neuron
                    //         // Synapse synapseForUpdate = this.currentNucleus.GetSynapse(oldElementNeuron);
                    //         // synapseForUpdate.nucleus = newElementNeuron;
                    //     }
                    // }
                    // else {
                    // it is a neuron in a subcluster
                    synapseNeuron.RemoveReceiver(this.currentNucleus);
                    newNucleus.AddReceiver(this.currentNucleus);
                    // }
                }
                else {
                    synapseNeuron.RemoveReceiver(this.currentNucleus);
                    newNucleus.AddReceiver(this.currentNucleus);
                }
            }

            protected virtual void DisconnectNucleus(Neuron nucleus) {
                if (this.currentNucleus.clusterPrefab == null)
                    return;
                string[] names = this.currentNucleus.synapses.Select(synapse => synapse.neuron.name).ToArray();
                int selectedIndex = -1;
                selectedIndex = EditorGUILayout.Popup("Disconnect from", selectedIndex, names);
                if (selectedIndex >= 0 && selectedIndex < this.currentNucleus.clusterPrefab.nuclei.Count) {
                    Synapse synapse = this.currentNucleus.synapses[selectedIndex];
                    Neuron synapseNeuron = synapse.neuron as Neuron;
                    synapseNeuron.RemoveReceiver(this.currentNucleus);
                }
            }

            #endregion Synapses 

            #endregion Inspector
        }
    }

}