From 5ed3e380c42f6d3cdb7c946eb64b191b888cc93a Mon Sep 17 00:00:00 2001
From: Andrettin <andre.ng@live.com>
Date: Thu, 7 Dec 2017 18:52:50 +0100
Subject: [PATCH] Improved worker automatic level-up upgrade choice

---
 src/action/action_resource.cpp |  6 ------
 src/include/unittype.h         |  1 +
 src/unit/unit.cpp              | 17 +++++++++++------
 src/unit/unittype.cpp          | 35 ++++++++++++++++++++++++++++++++++
 4 files changed, 47 insertions(+), 12 deletions(-)

diff --git a/src/action/action_resource.cpp b/src/action/action_resource.cpp
index b509daef8..dba0e485b 100644
--- a/src/action/action_resource.cpp
+++ b/src/action/action_resource.cpp
@@ -977,14 +977,8 @@ int COrder_Resource::GatherResource(CUnit &unit)
 		}
 
 		// Calculate how much we can load.
-		//Wyrmgus start
-//		if (resinfo.ResourceStep) {
 		if (unit.GetResourceStep(this->CurrentResource) && (harvest_from_outside || Map.Info.IsPointOnMap(this->goalPos, this->MapLayer))) {
-		//Wyrmgus end
-			//Wyrmgus start
-//			addload = resinfo.ResourceStep;
 			addload = unit.GetResourceStep(this->CurrentResource);
-			//Wyrmgus end
 		} else {
 			//Wyrmgus start
 //			addload = resinfo.ResourceCapacity;
diff --git a/src/include/unittype.h b/src/include/unittype.h
index 525bae454..e60583f40 100644
--- a/src/include/unittype.h
+++ b/src/include/unittype.h
@@ -905,6 +905,7 @@ class CUnitType
 	//Wyrmgus start
 	void RemoveButtons(int button_action = -1, std::string mod_file = "");
 	int GetAvailableLevelUpUpgrades() const;
+	int GetResourceStep(const int resource, const int player) const;
 	VariationInfo *GetDefaultVariation(CPlayer &player, int image_layer = -1) const;
 	VariationInfo *GetVariation(std::string variation_name, int image_layer = -1) const;
 	std::string GetRandomVariationIdent(int image_layer = -1) const;
diff --git a/src/unit/unit.cpp b/src/unit/unit.cpp
index 600648c81..1fcbac5ce 100644
--- a/src/unit/unit.cpp
+++ b/src/unit/unit.cpp
@@ -715,14 +715,19 @@ void CUnit::IncreaseLevel(int level_quantity, bool automatic_learning)
 				//if is a harvester who is currently gathering, try to upgrade to a unit type which is best for harvesting the current resource
 				unsigned int best_gathering_rate = 0;
 				for (size_t i = 0; i != AiHelpers.ExperienceUpgrades[Type->Slot].size(); ++i) {
-					if (CheckDependByType(*this, *AiHelpers.ExperienceUpgrades[Type->Slot][i], true)) {
-						if (Character == NULL || !Character->ForbiddenUpgrades[AiHelpers.ExperienceUpgrades[Type->Slot][i]->Slot]) {
-							if (AiHelpers.ExperienceUpgrades[Type->Slot][i]->ResInfo[this->CurrentResource] && AiHelpers.ExperienceUpgrades[Type->Slot][i]->ResInfo[this->CurrentResource]->ResourceStep >= best_gathering_rate) {
-								if (AiHelpers.ExperienceUpgrades[Type->Slot][i]->ResInfo[this->CurrentResource]->ResourceStep > best_gathering_rate) {
-									best_gathering_rate = AiHelpers.ExperienceUpgrades[Type->Slot][i]->ResInfo[this->CurrentResource]->ResourceStep;
+					CUnitType *experience_upgrade_type = AiHelpers.ExperienceUpgrades[Type->Slot][i];
+					if (CheckDependByType(*this, *experience_upgrade_type, true)) {
+						if (Character == NULL || !Character->ForbiddenUpgrades[experience_upgrade_type->Slot]) {
+							if (!experience_upgrade_type->ResInfo[this->CurrentResource]) {
+								continue;
+							}
+							int gathering_rate = experience_upgrade_type->GetResourceStep(this->CurrentResource, this->Player->Index);
+							if (gathering_rate >= best_gathering_rate) {
+								if (gathering_rate > best_gathering_rate) {
+									best_gathering_rate = gathering_rate;
 									potential_upgrades.clear();
 								}
-								potential_upgrades.push_back(AiHelpers.ExperienceUpgrades[Type->Slot][i]);
+								potential_upgrades.push_back(experience_upgrade_type);
 							}
 						}
 					}
diff --git a/src/unit/unittype.cpp b/src/unit/unittype.cpp
index f7a46ab2d..a168b6c8b 100644
--- a/src/unit/unittype.cpp
+++ b/src/unit/unittype.cpp
@@ -797,6 +797,41 @@ int CUnitType::GetAvailableLevelUpUpgrades() const
 	return value;
 }
 
+int CUnitType::GetResourceStep(const int resource, const int player) const
+{
+	if (!this->ResInfo[resource]) {
+		return 0;
+	}
+
+	int resource_step = this->ResInfo[resource]->ResourceStep;
+	
+	resource_step += this->Stats[player].Variables[GATHERINGBONUS_INDEX].Value;
+	
+	if (resource == CopperCost) {
+		resource_step += this->Stats[player].Variables[COPPERGATHERINGBONUS_INDEX].Value;
+	} else if (resource == SilverCost) {
+		resource_step += this->Stats[player].Variables[SILVERGATHERINGBONUS_INDEX].Value;
+	} else if (resource == GoldCost) {
+		resource_step += this->Stats[player].Variables[GOLDGATHERINGBONUS_INDEX].Value;
+	} else if (resource == WoodCost) {
+		resource_step += this->Stats[player].Variables[LUMBERGATHERINGBONUS_INDEX].Value;
+	} else if (resource == StoneCost || resource == LimestoneCost) {
+		resource_step += this->Stats[player].Variables[STONEGATHERINGBONUS_INDEX].Value;
+	} else if (resource == CoalCost) {
+		resource_step += this->Stats[player].Variables[COALGATHERINGBONUS_INDEX].Value;
+	} else if (resource == JewelryCost) {
+		resource_step += this->Stats[player].Variables[JEWELRYGATHERINGBONUS_INDEX].Value;
+	} else if (resource == FurnitureCost) {
+		resource_step += this->Stats[player].Variables[FURNITUREGATHERINGBONUS_INDEX].Value;
+	} else if (resource == LeatherCost) {
+		resource_step += this->Stats[player].Variables[LEATHERGATHERINGBONUS_INDEX].Value;
+	} else if (resource == DiamondsCost || resource == EmeraldsCost) {
+		resource_step += this->Stats[player].Variables[GEMSGATHERINGBONUS_INDEX].Value;
+	}
+	
+	return resource_step;
+}
+
 VariationInfo *CUnitType::GetDefaultVariation(CPlayer &player, int image_layer) const
 {
 	int variation_max = image_layer == -1 ? VariationMax : this->LayerVarInfo[image_layer].size();