Fix spells banned on map appearing in towns, cleanup code

lib/gameState/CGameState.cpp

@@ -884,41 +884,45 @@ void CGameState::initTowns()
 		//init spells
 		vti->spells.resize(GameConstants::SPELL_LEVELS);
 		vti->possibleSpells -= SpellID::PRESET;
+
 		for(ui32 z=0; z<vti->obligatorySpells.size();z++)
 		{
 			const auto * s = vti->obligatorySpells[z].toSpell();
 			vti->spells[s->getLevel()-1].push_back(s->id);
 			vti->possibleSpells -= s->id;
 		}
-		while(!vti->possibleSpells.empty())
+
+		vstd::erase_if(vti->possibleSpells, [&](const SpellID & spellID)
 		{
-			ui32 total=0;
-			int sel = -1;
+			const auto * spell = spellID.toSpell();
 
-			for(ui32 ps=0;ps<vti->possibleSpells.size();ps++)
-				total += vti->possibleSpells[ps].toSpell()->getProbability(vti->getFactionID());
+			if (spell->getProbability(vti->getFactionID()) == 0)
+				return true;
 
-			if (total == 0) // remaining spells have 0 probability
-				break;
+			if (spell->isSpecial() || spell->isCreatureAbility())
+				return true;
+
+			if (!isAllowed(spellID))
+				return true;
+
+			return false;
+		});
+
+		std::vector<int> spellWeights;
+		for (auto & spellID : vti->possibleSpells)
+			spellWeights.push_back(spellID.toSpell()->getProbability(vti->getFactionID()));
 
-			auto r = getRandomGenerator().nextInt(total - 1);
-			for(ui32 ps=0; ps<vti->possibleSpells.size();ps++)
-			{
-				r -= vti->possibleSpells[ps].toSpell()->getProbability(vti->getFactionID());
-				if(r<0)
-				{
-					sel = ps;
-					break;
-				}
-			}
-			if(sel<0)
-				sel=0;
 
-			const auto * s = vti->possibleSpells[sel].toSpell();
+		while(!vti->possibleSpells.empty())
+		{
+			size_t index = RandomGeneratorUtil::nextItemWeighted(spellWeights, getRandomGenerator());
+
+			const auto * s = vti->possibleSpells[index].toSpell();
 			vti->spells[s->getLevel()-1].push_back(s->id);
-			vti->possibleSpells -= s->id;
+
+			vti->possibleSpells.erase(vti->possibleSpells.begin() + index);
+			spellWeights.erase(spellWeights.begin() + index);
 		}
-		vti->possibleSpells.clear();
 	}
 }