More fixes to per-stage mass and cost values
More fixes to per-stage mass and cost values

--- a/KerbalEngineer/VesselSimulator/PartSim.cs
+++ b/KerbalEngineer/VesselSimulator/PartSim.cs
@@ -38,7 +38,7 @@
         public double baseMass;
         public double baseMassForCoM;
         public Vector3d centerOfMass;
-        public double cost;
+        public double baseCost;
         public int decoupledInStage;
         public bool fuelCrossFeed;
         public List<PartSim> fuelTargets = new List<PartSim>();
@@ -90,6 +90,7 @@
             partSim.resourceDrains.Reset();
             partSim.resourceFlowStates.Reset();
             partSim.resources.Reset();
+            partSim.baseCost = 0d;
             partSim.baseMass = 0d;
             partSim.baseMassForCoM = 0d;
             partSim.startMass = 0d;
@@ -122,7 +123,7 @@
             partSim.inverseStage = partSim.part.inverseStage;
             //MonoBehaviour.print("inverseStage = " + inverseStage);
 
-            partSim.cost = partSim.part.GetCostWet();
+            partSim.baseCost = partSim.part.GetCostDry();
 
             if (log != null)
             {
@@ -146,6 +147,8 @@
                 partSim.realMass = partSim.part.mass;
                 if (log != null) log.buf.AppendLine("Using part.mass of " + partSim.part.mass);
             }
+
+            partSim.moduleMass = partSim.part.GetModuleMass((float)partSim.realMass);
 
             for (int i = 0; i < partSim.part.Resources.Count; i++)
             {
@@ -408,6 +411,9 @@
 
         public double GetMass(int currentStage, bool forCoM = false)
         {
+            if (decoupledInStage >= currentStage)
+                return 0d;
+
             double mass = forCoM ? baseMassForCoM : baseMass;
 
             for (int i = 0; i < resources.Types.Count; ++i)
@@ -422,7 +428,22 @@
 
             return mass;
         }
-                
+
+        public double GetCost(int currentStage)
+        {
+            if (decoupledInStage >= currentStage)
+                return 0d;
+
+            double cost = baseCost;
+
+            for (int i = 0; i < resources.Types.Count; ++i)
+            {
+                cost += resources.GetResourceCost(resources.Types[i]);
+            }
+
+            return cost;
+        }
+
         public void ReleasePart()
         {
             this.part = null;

--- a/KerbalEngineer/VesselSimulator/ResourceContainer.cs
+++ b/KerbalEngineer/VesselSimulator/ResourceContainer.cs
@@ -150,6 +150,12 @@
             return density == 0d ? 0d : this.resources[type] * density;
         }
 
+        public double GetResourceCost(int type)
+        {
+            double unitCost = GetResourceUnitCost(type);
+            return unitCost == 0d ? 0d : this.resources[type] * unitCost;
+        }
+
         public static ResourceFlowMode GetResourceFlowMode(int type)
         {
             return PartResourceLibrary.Instance.GetDefinition(type).resourceFlowMode;
@@ -157,7 +163,7 @@
 
         public static ResourceTransferMode GetResourceTransferMode(int type)
         {
-            return PartResourceLibrary.Instance.GetDefinition(type).resourceTransferMode;;
+            return PartResourceLibrary.Instance.GetDefinition(type).resourceTransferMode;
         }
 
         public static float GetResourceDensity(int type)
@@ -169,5 +175,10 @@
         {
             return PartResourceLibrary.Instance.GetDefinition(type).name;
         }
+
+        public static double GetResourceUnitCost(int type)
+        {
+            return PartResourceLibrary.Instance.GetDefinition(type).unitCost;
+        }
     }
 }

--- a/KerbalEngineer/VesselSimulator/Simulation.cs
+++ b/KerbalEngineer/VesselSimulator/Simulation.cs
@@ -384,34 +384,34 @@
             // Create the array of stages that will be returned
             Stage[] stages = new Stage[this.currentStage + 1];
 
+            int startStage = currentStage;
+
             // Loop through the stages
             while (this.currentStage >= 0)
             {
                 if (log != null)
                 {
                     log.buf.AppendLine("Simulating stage " + this.currentStage);
-                    log.buf.AppendLine("ShipMass = " + this.ShipMass);
                     log.Flush();
                     this._timer.Reset();
                     this._timer.Start();
                 }
 
+                // Update active engines and resource drains
+                this.UpdateResourceDrains();
+
                 // Update the masses of the parts to correctly handle "no physics" parts
-                this.UpdatePartMasses();
+                this.stageStartMass = this.UpdatePartMasses();
 
                 if (log != null)
                     this.allParts[0].DumpPartToBuffer(log.buf, "", this.allParts);
 
-                // Update active engines and resource drains
-                this.UpdateResourceDrains();
-
                 // Create the Stage object for this stage
                 Stage stage = new Stage();
 
                 this.stageTime = 0d;
                 this.vecStageDeltaV = Vector3.zero;
 
-                this.stageStartMass = this.ShipMass;
                 this.stageStartCom = this.ShipCom;
 
                 this.stepStartMass = this.stageStartMass;
@@ -450,22 +450,31 @@
 
                 stage.thrustOffsetAngle = Math.Asin(sinThrustOffsetAngle) * 180 / Math.PI;
 
-                // Calculate the cost and mass of this stage and add all engines and tanks that are decoupled
-                // in the next stage to the dontStageParts list
+                // Calculate the total cost of the vessel at this point
+                stage.totalCost = 0d;
                 for (int i = 0; i < allParts.Count; ++i)
                 {
-                    PartSim partSim = allParts[i];
-
-                    if (partSim.decoupledInStage == this.currentStage - 1)
-                    {
-                        stage.cost += partSim.cost;
-                        stage.mass += partSim.GetStartMass();
-                    }
-
-                    if (partSim.hasVessel == false && partSim.isFairing && partSim.inverseStage == currentStage)
-                    {
-                        stage.mass += partSim.moduleMass;
-                    }
+                    if (this.currentStage > allParts[i].decoupledInStage)
+                        stage.totalCost += allParts[i].GetCost(currentStage);
+                }
+
+                // The total mass is simply the mass at the start of the stage
+                stage.totalMass = this.stageStartMass;
+
+                // If we have done a previous stage
+                if (currentStage < startStage)
+                {
+                    // Calculate what the previous stage's mass and cost were by subtraction
+                    Stage prev = stages[currentStage + 1];
+                    prev.cost = prev.totalCost - stage.totalCost;
+                    prev.mass = prev.totalMass - stage.totalMass;
+                }
+
+                // The above code will never run for the last stage so set those directly
+                if (currentStage == 0)
+                {
+                    stage.cost = stage.totalCost;
+                    stage.mass = stage.totalMass;
                 }
 
                 this.dontStageParts = dontStagePartsLists[this.currentStage];
@@ -628,8 +637,6 @@
                 // For each stage we total up the cost, mass, deltaV and time for this stage and all the stages above
                 for (int j = i; j >= 0; j--)
                 {
-                    stages[i].totalCost += stages[j].cost;
-                    stages[i].totalMass += stages[j].mass;
                     stages[i].totalDeltaV += stages[j].deltaV;
                     stages[i].totalTime += stages[j].time;
                     stages[i].partCount = i > 0 ? stages[i].totalPartCount - stages[i - 1].totalPartCount : stages[i].totalPartCount;
@@ -658,7 +665,7 @@
             return stages;
         }
 
-        public void UpdatePartMasses()
+        public double UpdatePartMasses()
         {
             for (int i = 0; i < this.allParts.Count; i++)
             {
@@ -688,8 +695,13 @@
                 }
             }
 
+            double totalMass = 0d;
             for (int i = 0; i < this.allParts.Count; i++)
-                this.allParts[i].startMass = this.allParts[i].GetMass(currentStage);
+            {
+                totalMass += this.allParts[i].startMass = this.allParts[i].GetMass(currentStage);
+            }
+
+            return totalMass;
         }
 
         // Make sure we free them all, even if they should all be free already at this point