--- a/FrmArmyTree.cs	Sat Mar 17 21:03:33 2012 +0000
+++ b/FrmArmyTree.cs	Sat Apr 14 20:48:13 2012 +0100
@@ -298,9 +298,24 @@
 		private void treeView_AfterSelect(object sender, TreeViewEventArgs e)
 		{
 			TreeNode node = treeView.SelectedNode;
-			bttnAdd.Enabled = (node != null && node.Tag is ArmyCategory);
-			bttnDelete.Enabled = (node != null && node.Tag is Unit);
-			bttnEdit.Enabled = (node != null && (node.Tag is Unit || node.Tag is Army));
+			Object tag = (node != null) ? node.Tag : null;
+			bool tagIsUnit = (tag is Unit);
+			bool tagIsUnitWithContained;
+
+			if (tagIsUnit)
+			{
+				tagIsUnitWithContained = ((Unit)tag).UnitType.ContainedUnitTypes.Length > 0;
+			}
+			else
+			{
+				tagIsUnitWithContained = false;
+			}
+
+			bool tagIsArmy = (tag is Army);
+			bool tagIsCategory = (tag is ArmyCategory);
+			bttnAdd.Enabled = tagIsCategory || tagIsUnitWithContained;
+			bttnDelete.Enabled = tagIsUnit;
+			bttnEdit.Enabled = tagIsArmy || tagIsUnit;
 		}
 
 		private void miEditArmy_Click(object sender, EventArgs e)
@@ -330,6 +345,17 @@
 			TreeNode node = treeView.SelectedNode;
 			Object tag = (node != null) ? node.Tag : null;
 			bool tagIsUnit = (tag is Unit);
+			bool tagIsUnitWithContained;
+
+			if (tagIsUnit)
+			{
+				tagIsUnitWithContained = ((Unit)tag).UnitType.ContainedUnitTypes.Length > 0;
+			}
+			else
+			{
+				tagIsUnitWithContained = false;
+			}
+
 			bool tagIsArmy = (tag is Army);
 			bool tagIsCategory = (tag is ArmyCategory);
 
@@ -342,7 +368,7 @@
 				miDeleteUnit.Visible = tagIsUnit;
 				miEditUnit.Visible = tagIsUnit;
 				miEditArmy.Visible = tagIsArmy;
-				miAddUnit.Visible = tagIsCategory;
+				miAddUnit.Visible = tagIsCategory || tagIsUnitWithContained;
 			}
 		}
 
@@ -353,8 +379,28 @@
 
 		private void AddUnit()
 		{
-			ArmyCategory tagData = treeView.SelectedNode.Tag as ArmyCategory;
-			AddUnitFromCategory(tagData);
+			object tag = treeView.SelectedNode.Tag;
+
+			if (tag is Unit)
+			{
+				Unit unit = (Unit)tag;
+				UnitType[] unitTypes = unit.UnitType.ContainedUnitTypes;
+				FrmNewUnit newUnitFrm = new FrmNewUnit(WarFoundryCore.CurrentArmy, unitTypes);
+				DialogResult dr = newUnitFrm.ShowDialog(this);
+
+				if (dr == DialogResult.OK)
+				{
+					UnitType newUnit = newUnitFrm.SelectedUnit;
+					CreateAndAddUnitCommand cmd = new CreateAndAddUnitCommand(newUnitFrm.SelectedUnit, WarFoundryCore.CurrentArmy.GetCategory(newUnit.MainCategory));
+					commandStack.Execute(cmd);
+					MdiParent.OpenUnitDialog(cmd.Unit);
+				}
+			}
+			else
+			{
+				ArmyCategory tagData = tag as ArmyCategory;
+				AddUnitFromCategory(tagData);
+			}
 		}
 
 		private void AddUnitFromCategory(ArmyCategory tagData)