From e7087fa6020565b322acc2dd69df9cd3f874ac1c Mon Sep 17 00:00:00 2001
From: JoJo <joseph.salomone@dealertrack.com>
Date: Thu, 30 Dec 2021 12:57:54 -0800
Subject: [PATCH] Some minor fixes

---
 .../Builder/General/HeavenlyBodyBuilder.cs    | 22 +++++++++++++------
 .../Builder/General/PlanetDestroyer.cs        |  2 ++
 NewHorizons/Main.cs                           |  2 +-
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/NewHorizons/Builder/General/HeavenlyBodyBuilder.cs b/NewHorizons/Builder/General/HeavenlyBodyBuilder.cs
index 779156a1..036ca37c 100644
--- a/NewHorizons/Builder/General/HeavenlyBodyBuilder.cs
+++ b/NewHorizons/Builder/General/HeavenlyBodyBuilder.cs
@@ -38,18 +38,18 @@ namespace NewHorizons.Builder.General
                 hb = AddHeavenlyBody(config.Name);
             }
 
-            if(!config.Orbit.IsStatic)
+            var mapping = Planet.defaultMapping;
+            if (!config.Orbit.IsStatic)
             {
                 var planetoid = new Planet.Plantoid(size, gravity, body.transform.rotation, initialMotion._initAngularSpeed, parent, orbit);
 
-                var mapping = Planet.defaultMapping;
                 mapping[hb] = planetoid;
                 Planet.defaultMapping = mapping;
             }
 
             // Fix for binary focal points
             var focalPoint = Position.AstroLookup[parent].Invoke()?.gameObject.GetComponent<BinaryFocalPoint>();
-            if (focalPoint != null && Planet.defaultMapping.ContainsKey(parent))
+            if (focalPoint != null && mapping.ContainsKey(parent))
             {
                 var primary = Position.getBody(GetBody(focalPoint.PrimaryName));
                 var secondary = Position.getBody(GetBody(focalPoint.SecondaryName));
@@ -63,13 +63,21 @@ namespace NewHorizons.Builder.General
                     var exponent = (primaryGravity.exponent + secondaryGravity.exponent) / 2f;
                     var mass = (primaryGravity.mass + secondaryGravity.mass) / 4f;
 
-                    var currentValue = Planet.mapping[parent];
-                    var newValue = new Planet.Plantoid(currentValue.size, Gravity.of(exponent, mass), currentValue.state);
-                    Planet.defaultMapping[parent] = newValue;
+                    mapping[parent] = new Planet.Plantoid(mapping[parent].size, Gravity.of(exponent, mass), mapping[parent].state);
+                    Planet.defaultMapping = mapping;
                 }
             }
         }
 
+        public static void Remove(AstroObject obj)
+        {
+            var astro = Position.find(obj);
+
+            var mapping = Planet.defaultMapping;
+            mapping.Remove(astro);
+            Planet.defaultMapping = mapping;
+        }
+
         private static Gravity getGravity(GravityVolume volume)
         {
             if (volume == null)
@@ -115,7 +123,7 @@ namespace NewHorizons.Builder.General
             return hb;
         }
 
-        public static void OnDestroy()
+        public static void Reset()
         {
             Planet.defaultMapping = Planet.standardMapping;
         }
diff --git a/NewHorizons/Builder/General/PlanetDestroyer.cs b/NewHorizons/Builder/General/PlanetDestroyer.cs
index 35740cae..6d61c057 100644
--- a/NewHorizons/Builder/General/PlanetDestroyer.cs
+++ b/NewHorizons/Builder/General/PlanetDestroyer.cs
@@ -102,6 +102,8 @@ namespace NewHorizons.Builder.General
             Main.Instance.ModHelper.Events.Unity.FireOnNextUpdate(() => RemoveProxy(ao.name.Replace("_Body", "")));
 
             ao.transform.root.gameObject.SetActive(false);
+
+            HeavenlyBodyBuilder.Remove(ao);
         }
 
         public static void RemoveDistantProxyClones()
diff --git a/NewHorizons/Main.cs b/NewHorizons/Main.cs
index b0abf0fe..4ce1df2e 100644
--- a/NewHorizons/Main.cs
+++ b/NewHorizons/Main.cs
@@ -70,7 +70,7 @@ namespace NewHorizons
         void OnSceneLoaded(Scene scene, LoadSceneMode mode)
         {
             Logger.Log($"Scene Loaded: {scene.name} {mode}");
-            HeavenlyBodyBuilder.OnDestroy();
+            HeavenlyBodyBuilder.Reset();
             if (scene.name != "SolarSystem") { return; }
 
             Instance.ModHelper.Events.Unity.FireOnNextUpdate(() => Locator.GetPlayerBody().gameObject.AddComponent<DebugRaycaster>());