From dcf845b583038f0cc9fc8427fb41a7352cc681b0 Mon Sep 17 00:00:00 2001
From: Alexey Chernyshov <65975574+artifixer@users.noreply.github.com>
Date: Tue, 14 May 2024 02:27:29 +0300
Subject: [PATCH] Updates to DistanceMatrix

---
 build/common.props                            |  2 +-
 changelog.txt                                 |  4 ++
 .../DistanceMatrixImplementation.cs           | 43 ++++++++++++++++---
 .../SubModule.cs                              |  3 ++
 .../DistanceMatrix/DistanceMatrixT.cs         |  7 +++
 5 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/build/common.props b/build/common.props
index 3123d4f3..2951849e 100644
--- a/build/common.props
+++ b/build/common.props
@@ -5,7 +5,7 @@
   <PropertyGroup>
     <BaseGameVersion>1.0.0</BaseGameVersion>
     <!--Module Version-->
-    <Version>2.9.7</Version>
+    <Version>2.9.8</Version>
     <!--Harmony Version-->
     <HarmonyVersion>2.2.2</HarmonyVersion>
     <HarmonyExtensionsVersion>3.2.0.77</HarmonyExtensionsVersion>
diff --git a/changelog.txt b/changelog.txt
index d2d43d4f..17813418 100644
--- a/changelog.txt
+++ b/changelog.txt
@@ -1,4 +1,8 @@
 ---------------------------------------------------------------------------------------------------
+Version: 2.9.8
+Game Versions: v1.0.x,v1.1.x,v1.2.x
+* Fixed DistanceMatrix
+---------------------------------------------------------------------------------------------------
 Version: 2.9.7
 Game Versions: v1.0.0,v1.0.1,v1.0.2,v1.0.3,v1.1.0,v1.1.1,v1.1.2,v1.1.3,v1.1.4,v1.1.5,v1.1.6,v1.2.x
 * Fixed UpdateInfo parsing for Crash Reports
diff --git a/src/Bannerlord.ButterLib.Implementation/DistanceMatrix/DistanceMatrixImplementation.cs b/src/Bannerlord.ButterLib.Implementation/DistanceMatrix/DistanceMatrixImplementation.cs
index c76b9ab8..71d56c71 100644
--- a/src/Bannerlord.ButterLib.Implementation/DistanceMatrix/DistanceMatrixImplementation.cs
+++ b/src/Bannerlord.ButterLib.Implementation/DistanceMatrix/DistanceMatrixImplementation.cs
@@ -17,6 +17,8 @@ internal sealed class DistanceMatrixImplementation<T> : DistanceMatrix<T> where
     //Fields
     private readonly Dictionary<ulong, float> _distanceMatrix;
     private readonly Dictionary<(T Object1, T Object2), float> _typedDistanceMatrix;
+    private readonly Dictionary<T, SortedSet<(T OtherObject, float Distance)>> _flatenedDictionary;
+
     private readonly Func<IEnumerable<T>>? _entityListGetter;
     private readonly Func<T, T, float>? _distanceCalculator;
     private Dictionary<MBGUID, MBObjectBase> _cachedMapping = new();
@@ -28,6 +30,8 @@ internal sealed class DistanceMatrixImplementation<T> : DistanceMatrix<T> where
     /// <inheritdoc/>
     public override Dictionary<(T Object1, T Object2), float> AsTypedDictionary => _typedDistanceMatrix;
 
+    public override Dictionary<T, SortedSet<(T OtherObject, float Distance)>> AsFlatenedDictionary => _flatenedDictionary;
+
     //Constructors
     /// <inheritdoc/>
     public DistanceMatrixImplementation()
@@ -36,6 +40,7 @@ public DistanceMatrixImplementation()
         _distanceCalculator = null;
         _distanceMatrix = CalculateDistanceMatrix();
         _typedDistanceMatrix = GetTypedDistanceMatrix();
+        _flatenedDictionary = GetFlatenedDictionary();
     }
 
     /// <inheritdoc/>
@@ -45,6 +50,7 @@ public DistanceMatrixImplementation(Func<IEnumerable<T>> customListGetter, Func<
         _distanceCalculator = customDistanceCalculator;
         _distanceMatrix = CalculateDistanceMatrix();
         _typedDistanceMatrix = GetTypedDistanceMatrix();
+        _flatenedDictionary = GetFlatenedDictionary();
     }
 
     //Public methods
@@ -58,6 +64,11 @@ public override void SetDistance(T object1, T object2, float distance)
     {
         _distanceMatrix[object1.Id > object2.Id ? ElegantPairHelper.Pair(object2.Id, object1.Id) : ElegantPairHelper.Pair(object1.Id, object2.Id)] = distance;
         _typedDistanceMatrix[object1.Id > object2.Id ? (object2, object1) : (object1, object2)] = distance;
+        
+        _flatenedDictionary[object1].RemoveWhere(x => x.OtherObject == object2);
+        _flatenedDictionary[object1].Add((object2, distance));
+        _flatenedDictionary[object2].RemoveWhere(x => x.OtherObject == object1);
+        _flatenedDictionary[object2].Add((object1, distance));
     }
 
     //Private methods
@@ -89,7 +100,7 @@ private Dictionary<ulong, float> CalculateDistanceMatrix()
         if (typeof(Hero).IsAssignableFrom(typeof(T)))
         {
             var activeHeroes = Hero.AllAliveHeroes
-                .Where(h => h.IsInitialized && !h.IsNotSpawned && !h.IsDisabled && !h.IsDead && !h.IsNotable).ToList();
+                .Where(h => !h.IsNotSpawned && !h.IsDisabled && !h.IsDead && !h.IsChild && !h.IsNotable).ToList();
             _cachedMapping = activeHeroes.ToDictionary(key => key.Id, value => value as MBObjectBase);
 
             return activeHeroes
@@ -102,7 +113,7 @@ private Dictionary<ulong, float> CalculateDistanceMatrix()
 
         if (typeof(Settlement).IsAssignableFrom(typeof(T)))
         {
-            var settlements = Settlement.All.Where(s => s.IsInitialized && (s.IsFortification || s.IsVillage)).ToList();
+            var settlements = Settlement.All.Where(s => s.IsFortification || s.IsVillage).ToList();
             _cachedMapping = settlements.ToDictionary(key => key.Id, value => value as MBObjectBase);
             return settlements
                 .SelectMany(_ => settlements, (X, Y) => (X, Y))
@@ -111,10 +122,10 @@ private Dictionary<ulong, float> CalculateDistanceMatrix()
                     key => ElegantPairHelper.Pair(key.X.Id, key.Y.Id),
                     value => Campaign.Current.Models.MapDistanceModel.GetDistance(value.X, value.Y));
         }
-
+        
         if (typeof(Clan).IsAssignableFrom(typeof(T)))
         {
-            var clans = Clan.All.Where(c => c.IsInitialized && c.Fiefs.Any()).ToList();
+            var clans = Clan.All.Where(c => !c.IsEliminated && c.Fiefs.Any()).ToList();
             _cachedMapping = clans.ToDictionary(key => key.Id, value => value as MBObjectBase);
 
             var settlementDistanceMatrix = Campaign.Current.GetCampaignBehavior<GeopoliticsBehavior>().SettlementDistanceMatrix ?? new DistanceMatrixImplementation<Settlement>();
@@ -130,7 +141,7 @@ private Dictionary<ulong, float> CalculateDistanceMatrix()
 
         if (typeof(Kingdom).IsAssignableFrom(typeof(T)))
         {
-            var kingdoms = Kingdom.All.Where(k => k.IsInitialized && k.Fiefs.Any()).ToList();
+            var kingdoms = Kingdom.All.Where(k => !k.IsEliminated && k.Fiefs.Any()).ToList();
             _cachedMapping = kingdoms.ToDictionary(key => key.Id, value => value as MBObjectBase);
 
             var claDistanceMatrix = Campaign.Current.GetCampaignBehavior<GeopoliticsBehavior>().ClanDistanceMatrix ?? new DistanceMatrixImplementation<Clan>();
@@ -145,4 +156,26 @@ private Dictionary<ulong, float> CalculateDistanceMatrix()
     private Dictionary<(T Object1, T Object2), float> GetTypedDistanceMatrix() => _distanceMatrix
         .ToDictionary(key => ElegantPairHelper.UnPairMBGUID(key.Key), value => value.Value)
         .ToDictionary(key => (GetObject(key.Key.A), GetObject(key.Key.B)), value => value.Value);
+
+    private Dictionary<T, SortedSet<(T OtherObject, float Distance)>> GetFlatenedDictionary()
+    {
+        var list = _typedDistanceMatrix.ToList();
+        var keyList = list.SelectMany(kvp => new[] { kvp.Key.Object1, kvp.Key.Object2 }).Distinct().ToList();
+        var result = new Dictionary<T, SortedSet<(T OtherObject, float Distance)>>();
+        keyList.ForEach(key =>
+        {
+            var valueList = list.Where(kvp => kvp.Key.Object1 == key || kvp.Key.Object2 == key).Select(kvp => (OtherObject: kvp.Key.Object1 == key ? kvp.Key.Object2 : kvp.Key.Object1, Distance: kvp.Value)).Distinct().ToList();
+             SortedSet<(T OtherObject, float Distance)> valueSet = new(valueList, new TupleComparer());
+            result.Add(key, valueSet);
+        });
+        return result;
+    }
+
+    private class TupleComparer : IComparer<(T OtherObject, float Distance)>
+    {
+        public int Compare((T OtherObject, float Distance) x, (T OtherObject, float Distance) y)
+        {
+            return Comparer<float>.Default.Compare(x.Distance, y.Distance);
+        }
+    }
 }
\ No newline at end of file
diff --git a/src/Bannerlord.ButterLib.Implementation/SubModule.cs b/src/Bannerlord.ButterLib.Implementation/SubModule.cs
index bc0b9a66..907bafab 100644
--- a/src/Bannerlord.ButterLib.Implementation/SubModule.cs
+++ b/src/Bannerlord.ButterLib.Implementation/SubModule.cs
@@ -1,6 +1,8 @@
 using Bannerlord.ButterLib.Common.Extensions;
 using Bannerlord.ButterLib.DistanceMatrix;
+using Bannerlord.ButterLib.Extensions;
 using Bannerlord.ButterLib.HotKeys;
+using Bannerlord.ButterLib.Implementation.Common.Extensions;
 using Bannerlord.ButterLib.Implementation.DistanceMatrix;
 using Bannerlord.ButterLib.Implementation.HotKeys;
 using Bannerlord.ButterLib.Implementation.Logging;
@@ -35,6 +37,7 @@ public void OnServiceRegistration()
         {
             services.AddScoped(typeof(DistanceMatrix<>), typeof(DistanceMatrixImplementation<>));
             services.AddSingleton<IDistanceMatrixStatic, DistanceMatrixStaticImplementation>();
+            services.AddSingleton<ICampaignExtensions, CampaignExtensionsImplementation>();
 
             services.AddScoped<IMBObjectExtensionDataStore, MBObjectExtensionDataStore>();
             services.AddScoped<IMBObjectFinder, MBObjectFinder>();
diff --git a/src/Bannerlord.ButterLib/DistanceMatrix/DistanceMatrixT.cs b/src/Bannerlord.ButterLib/DistanceMatrix/DistanceMatrixT.cs
index 68740403..69e496f5 100644
--- a/src/Bannerlord.ButterLib/DistanceMatrix/DistanceMatrixT.cs
+++ b/src/Bannerlord.ButterLib/DistanceMatrix/DistanceMatrixT.cs
@@ -77,6 +77,13 @@ protected DistanceMatrix(Func<IEnumerable<T>> customListGetter, Func<T, T, float
     /// </value>
     public abstract Dictionary<(T Object1, T Object2), float> AsTypedDictionary { get; }
 
+    /// <summary>Objectified distance matrix representation for nearest neighbours processing</summary>
+    /// <value>
+    /// A dictionary of the <typeparamref name="T"/> object type as a key
+    /// and a values are represented by a <see cref="T:System.Collections.Generic.SortedSet`1"/> containing tuples of all other objects and the corresponding distances in a form of floating point numbers.
+    /// </value>
+    public abstract Dictionary<T, SortedSet<(T OtherObject, float Distance)>> AsFlatenedDictionary { get; }
+
     /// <summary>Gets calculated distance between specified type <typeparamref name="T"/> objects.</summary>
     /// <param name="object1">The first of the objects between which it is necessary to determine the distance.</param>
     /// <param name="object2">The second of the objects between which it is necessary to determine the distance.</param>