Fix Grey Head Disease regression.

Sheep-y · Aug 28, 2018 · b623227 · b623227
1 parent 4c725db
commit b623227
Showing 1 changed file with 13 additions and 31 deletions.
diff --git a/AttackImprovementMod/src/modules/HitLocation.cs b/AttackImprovementMod/src/modules/HitLocation.cs
@@ -1,11 +1,7 @@
-using BattleTech.UI;
 using BattleTech;
 using System;
 using System.Collections.Generic;
-using System.Linq;
 using System.Reflection;
-using UnityEngine.EventSystems;
-using UnityEngine;
 
 namespace Sheepy.BattleTechMod.AttackImprovementMod {
    using static ArmorLocation;
@@ -44,7 +40,7 @@ public override void CombatStartsOnce () {
 
          if ( Settings.FixGreyHeadDisease ) {
             HeadHitWeights = new Dictionary<Dictionary<ArmorLocation, int>, int>();
-            Patch( MechGetHit, null, "FixGreyHeadDisease" );
+            Patch( typeof( BattleTech.HitLocation ), "GetMechHitTable", null, "FixGreyHeadDisease" );
          }
       }
 
@@ -73,22 +69,6 @@ public override void CombatStarts () {
                   }
             }
          }
-
-         if ( Settings.FixGreyHeadDisease ) {
-            List<Dictionary<ArmorLocation, int>> hitTables;
-            if ( Settings.FixHitDistribution ) 
-               hitTables = ScaledMechHitTables.Values.ToList();
-            else {
-               hitTables = new List<Dictionary<ArmorLocation, int>>();
-               foreach ( AttackDirection direction in Enum.GetValues( typeof( AttackDirection ) ) ) {
-                  if ( direction == AttackDirection.None ) continue;
-                  hitTables.Add( Combat.HitLocation.GetMechHitTable( direction ) );
-               }
-            }
-            foreach ( Dictionary<ArmorLocation, int> hitTable in hitTables )
-               if ( hitTable.TryGetValue( Head, out int head ) && head > 0 )
-                  HeadHitWeights.Add( hitTable, head );
-         }
       }
 
       public override void CombatEnds () {
@@ -138,12 +118,8 @@ public static void PrefixMechCalledShot ( ref Dictionary<ArmorLocation, int> hit
       }                 catch ( Exception ex ) { Error( ex ); } }
 
       public static void ScaleMechHitTable ( ref Dictionary<ArmorLocation, int> hitTable ) { try {
-         if ( ! ScaledMechHitTables.TryGetValue( hitTable, out Dictionary<ArmorLocation, int> scaled ) ) {
+         if ( ! ScaledMechHitTables.TryGetValue( hitTable, out Dictionary<ArmorLocation, int> scaled ) )
             ScaledMechHitTables.Add( hitTable, scaled = ScaleHitTable( hitTable ) );
-            Warn( "New unscaled hit table [{0}]", Join( ",", hitTable.Select( e => e.Key+"="+e.Value ) ) );
-            if ( Settings.FixGreyHeadDisease && scaled.TryGetValue( Head, out int head ) && head > 0 )
-               HeadHitWeights.Add( scaled, head ); // Would be too late if head is already removed, thus the warning
-         }
          hitTable = scaled;
       }                 catch ( Exception ex ) { Error( ex ); } }
 
@@ -174,11 +150,17 @@ internal static int SumWeight<T> ( Dictionary<T, int> hitTable, T bonusLocation,
          return totalWeight;
       }
 
-      // Not the most efficient fix since it is called per shot - same as the bugged head removal code - but this is dead simple
-      public static void FixGreyHeadDisease ( Dictionary<ArmorLocation, int> hitTable ) {
-         // Re-attach missing head after hit location is rolled
-         if ( ! hitTable.ContainsKey( Head ) && HeadHitWeights.ContainsKey( hitTable ) )
-            hitTable.Add( Head, HeadHitWeights[ hitTable ] );
+      public static void FixGreyHeadDisease ( Dictionary<ArmorLocation, int> __result ) {
+         if ( __result == null ) return;
+         if ( __result.TryGetValue( Head, out int head ) ) {
+            // Has head. Cache it.
+            if ( HeadHitWeights.ContainsKey( __result ) ) return;
+            HeadHitWeights[ __result ] = head;
+         } else {
+            // No head. Add it?
+            if ( ! HeadHitWeights.TryGetValue( __result, out head ) ) return;
+            __result[ Head ] = head;
+         }
       }
    }
 }