--- a/api/Objects/Army.cs	Tue Apr 07 14:43:04 2009 +0000
+++ b/api/Objects/Army.cs	Tue Apr 07 15:28:06 2009 +0000
@@ -3,8 +3,7 @@
 // The file and the library/program it is in are licensed under the GNU LGPL license, either version 3 of the License or (at your option) any later version. Please see COPYING.LGPL for more information and the full license.
 
 using System;
-using System.IO;
-using System.Collections;
+using System.IO;
 using System.Collections.Generic;
 using System.Text;
 using System.Xml;
@@ -24,7 +23,7 @@
 		private Race armyRace;
 		private int maxPoints;
 		private double pointsTotal;
-		private ArmyCategory[] categories;
+		private Dictionary<Category, ArmyCategory> categories;
 
 		public event ObjectAddDelegate UnitAdded;
 		public event ObjectRemoveDelegate UnitRemoved;
@@ -46,40 +45,42 @@
 		
 		public ArmyCategory GetCategory(Category cat)
 		{
-			foreach (ArmyCategory armyCat in Categories)
-			{
-				if (armyCat.Category.Equals(cat))
-				{
-					return armyCat;
-				}
-			}
-
-			return null;
+			ArmyCategory armyCat = null;
+			ArmyCategories.TryGetValue(cat, out armyCat);
+			return armyCat;
+		}
+		
+		private Dictionary<Category, ArmyCategory> ArmyCategories
+		{
+			get
+			{
+				if (categories==null)
+				{
+					categories = new Dictionary<Category, ArmyCategory>();
+					Category[] raceCats = Race.Categories;
+					ArmyCategory cat;
+					int raceCatCount = raceCats.Length;
+
+					for (int i = 0; i < raceCatCount; i++)
+					{
+						Category raceCat = raceCats[i];
+						cat = new ArmyCategory(this, raceCat);
+						categories[raceCat] = cat;
+						cat.PointsValueChanged+= PointsValueChangedMethod;
+						cat.UnitAdded+=new ObjectAddDelegate(Army_UnitAdded);
+						cat.UnitRemoved+=new ObjectRemoveDelegate(Army_UnitRemoved);
+					}
+				}
+				
+				return categories;
+			}
 		}
 
 		public ArmyCategory[] Categories
 		{
 			get 
 			{
-				if (categories==null)
-				{
-					Category[] raceCats = Race.Categories;
-					ArmyCategory cat;
-					int raceCatCount = raceCats.Length;
-					categories = new ArmyCategory[raceCatCount];
-
-					for (int i = 0; i < raceCatCount; i++)
-					{
-						cat = new ArmyCategory(this, raceCats[i]);
-						categories[i] = cat;
-						cat.PointsValueChanged+= PointsValueChangedMethod;
-						cat.UnitAdded+=new ObjectAddDelegate(Army_UnitAdded);
-						cat.UnitRemoved+=new ObjectRemoveDelegate(Army_UnitRemoved);
-						cat.RequirementsFailed+=new FailedUnitRequirementDelegate(Army_FailedRequirement);
-					}
-				}
-
-				return categories;
+				return DictionaryUtils.ToArray<Category, ArmyCategory>(ArmyCategories);
 			}
 		}
 
@@ -155,6 +156,24 @@
 		public double PointsTotal
 		{
 			get { return TotalPoints; }
+		}
+		
+		public void AddUnit(Unit unit)
+		{
+			List<FailedUnitRequirement> failedReqs = CanAddUnit(unit);
+			unit.Army = this;
+			ArmyCategory armyCat = GetCategory(unit.UnitType.MainCategory);
+			armyCat.AddUnit(unit);
+			OnUnitAdded(unit, failedReqs);
+		}
+		
+		public void RemoveUnit(Unit unit)
+		{
+			List<FailedUnitRequirement> failedReqs = CanRemoveUnit(unit);
+			unit.Army = null;
+			ArmyCategory armyCat = GetCategory(unit.UnitType.MainCategory);
+			armyCat.RemoveUnit(unit);
+			OnUnitRemoved(unit, failedReqs);
 		}
 
 		public Unit[] GetUnits(Category cat)
@@ -169,20 +188,14 @@
 
 		public Unit[] GetUnits()
 		{
-			ArrayList fullList = new ArrayList();
+			List<Unit> fullList = new List<Unit>();
 
 			foreach(ArmyCategory cat in Categories)
-			{
-				foreach(Unit unit in cat.GetUnits())
-				{
-					if (!fullList.Contains(unit))
-					{
-						fullList.Add(unit);
-					}
-				}
+			{
+				fullList.AddRange(cat.GetUnits());
 			}
 
-			return (Unit[])fullList.ToArray(typeof(Unit));
+			return fullList.ToArray();
 		}
 
 		public int MaxPoints
@@ -253,13 +266,5 @@
 		{
 			OnUnitRemoved((Unit)val);
 		}
-
-		private void Army_FailedRequirement(List<FailedUnitRequirement> failedRequirements)
-		{
-			if (FailedRequirement!=null)
-			{
-				FailedRequirement(failedRequirements);
-			}
-		}
 	}
 }