From 3c00f0b441daa783e0c1d659935d36c565f9d075 Mon Sep 17 00:00:00 2001
From: Andrettin <andre.ng@live.com>
Date: Thu, 17 Dec 2015 16:46:57 +0100
Subject: [PATCH] Some improvements to the hero code

---
 src/include/unittype.h |  1 +
 src/unit/unit.cpp      | 53 ++++++++++++++++++++++++++++++++++++++++--
 src/unit/unittype.cpp  | 31 ++++++++++++++++++++++++
 3 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/src/include/unittype.h b/src/include/unittype.h
index c99342758..bd6749115 100644
--- a/src/include/unittype.h
+++ b/src/include/unittype.h
@@ -656,6 +656,7 @@ class CUnitType
 	bool CanSelect(GroupSelectionMode mode = SELECTABLE_BY_RECTANGLE_ONLY) const;
 	
 	//Wyrmgus start
+	int GetAvailableLevelUpUpgrades() const;
 	VariationInfo *GetDefaultVariation(CPlayer &player) const;
 	std::string GetDefaultName(CPlayer &player) const;
 	CPlayerColorGraphic *GetDefaultLayerSprite(CPlayer &player, int image_layer) const;
diff --git a/src/unit/unit.cpp b/src/unit/unit.cpp
index c65c15c58..a717dcab4 100644
--- a/src/unit/unit.cpp
+++ b/src/unit/unit.cpp
@@ -58,6 +58,9 @@
 //Wyrmgus end
 #include "game.h"
 #include "editor.h"
+//Wyrmgus start
+#include "grand_strategy.h"
+//Wyrmgus end
 #include "interface.h"
 //Wyrmgus start
 #include "item.h"
@@ -663,8 +666,43 @@ void CUnit::Retrain()
 		}
 	}
 	
+	std::string unit_name = GetMessageName();
+	
+	//now, revert the unit's type to the level 1 one
+	while (this->Type->Stats[this->Player->Index].Variables[LEVEL_INDEX].Value > 1) {
+		bool found_previous_unit_type = false;
+		for (size_t i = 0; i != UnitTypes.size(); ++i) {
+			if (Character != NULL && Character->ForbiddenUpgrades[i]) {
+				continue;
+			}
+			for (size_t j = 0; j != AiHelpers.ExperienceUpgrades[i].size(); ++j) {
+				if (AiHelpers.ExperienceUpgrades[i][j] == this->Type) {
+					this->Variable[LEVELUP_INDEX].Value += 1;
+					TransformUnitIntoType(*this, *UnitTypes[i]);
+					if (!IsNetworkGame() && Character != NULL && Character->Persistent && Player->AiEnabled == false) {	//save the unit-type experience upgrade for persistent characters
+						if (Character->Type->Slot != i) {
+							Character->Type = const_cast<CUnitType *>(&(*UnitTypes[i]));
+							if (GrandStrategy) { //also update the corresponding grand strategy hero, if in grand strategy mode
+								CGrandStrategyHero *hero = GrandStrategyGame.GetHero(Character->GetFullName());
+								if (hero) {
+									hero->SetType(i);
+								}
+							}
+							SaveHero(Character);
+						}
+					}
+					found_previous_unit_type = true;
+					break;
+				}
+			}
+			if (found_previous_unit_type) {
+				break;
+			}
+		}
+	}
+	
 	if (this->Player == ThisPlayer) {
-		this->Player->Notify(NotifyGreen, this->tilePos, _("%s has retrained."), GetMessageName().c_str());
+		this->Player->Notify(NotifyGreen, this->tilePos, _("%s's level-up choices have been reset."), unit_name.c_str());
 	}
 }
 
@@ -3235,15 +3273,26 @@ PixelPos CUnit::GetMapPixelPosCenter() const
 int CUnit::GetAvailableLevelUpUpgrades(bool only_units) const
 {
 	int value = 0;
+	int upgrade_value = 0;
 	
 	if (AiHelpers.ExperienceUpgrades.size() > Type->Slot) {
 		for (size_t i = 0; i != AiHelpers.ExperienceUpgrades[Type->Slot].size(); ++i) {
 			if (Character == NULL || !Character->ForbiddenUpgrades[AiHelpers.ExperienceUpgrades[Type->Slot][i]->Slot]) {
-				value += 1;
+				int local_upgrade_value = 1;
+				
+				if (!only_units) {
+					local_upgrade_value += AiHelpers.ExperienceUpgrades[Type->Slot][i]->GetAvailableLevelUpUpgrades();
+				}
+				
+				if (local_upgrade_value > upgrade_value) {
+					upgrade_value = local_upgrade_value;
+				}
 			}
 		}
 	}
 	
+	value += upgrade_value;
+	
 	if (!only_units && AiHelpers.LearnableAbilities.size() > Type->Slot) {
 		for (size_t i = 0; i != AiHelpers.LearnableAbilities[Type->Slot].size(); ++i) {
 			if (!IndividualUpgrades[AiHelpers.LearnableAbilities[Type->Slot][i]->ID]) {
diff --git a/src/unit/unittype.cpp b/src/unit/unittype.cpp
index 726009970..fe620f968 100644
--- a/src/unit/unittype.cpp
+++ b/src/unit/unittype.cpp
@@ -37,6 +37,9 @@
 
 #include "unittype.h"
 
+//Wyrmgus start
+#include "../ai/ai_local.h" //for using AiHelpers
+//Wyrmgus end
 #include "animation.h"
 #include "animation/animation_exactframe.h"
 #include "animation/animation_frame.h"
@@ -726,6 +729,34 @@ bool CUnitType::CanSelect(GroupSelectionMode mode) const
 }
 
 //Wyrmgus start
+int CUnitType::GetAvailableLevelUpUpgrades() const
+{
+	int value = 0;
+	int upgrade_value = 0;
+	
+	if (AiHelpers.ExperienceUpgrades.size() > this->Slot) {
+		for (size_t i = 0; i != AiHelpers.ExperienceUpgrades[this->Slot].size(); ++i) {
+			int local_upgrade_value = 1;
+			
+			local_upgrade_value += AiHelpers.ExperienceUpgrades[this->Slot][i]->GetAvailableLevelUpUpgrades();
+			
+			if (local_upgrade_value > upgrade_value) {
+				upgrade_value = local_upgrade_value;
+			}
+		}
+	}
+	
+	value += upgrade_value;
+	
+	if (AiHelpers.LearnableAbilities.size() > this->Slot) {
+		for (size_t i = 0; i != AiHelpers.LearnableAbilities[this->Slot].size(); ++i) {
+			value += 1;
+		}
+	}
+	
+	return value;
+}
+
 VariationInfo *CUnitType::GetDefaultVariation(CPlayer &player) const
 {
 	for (int i = 0; i < VariationMax; ++i) {