Add category count

diff --git a/spec/public/app/code/local/Magehack/Elasticmage/Model/Resource/Product/CollectionSpec.php b/spec/public/app/code/local/Magehack/Elasticmage/Model/Resource/Product/CollectionSpec.php
index 666b20b..50d3369 100644
--- a/spec/public/app/code/local/Magehack/Elasticmage/Model/Resource/Product/CollectionSpec.php
+++ b/spec/public/app/code/local/Magehack/Elasticmage/Model/Resource/Product/CollectionSpec.php
@@ -2,6 +2,7 @@

 namespace spec;

+use PhpSpec\Exception\Exception;
 use PhpSpec\ObjectBehavior;
 use Prophecy\Argument;

@@ -390,4 +391,41 @@ class Magehack_Elasticmage_Model_Resource_Product_CollectionSpec extends ObjectB
         $this->_validateLoadedSampleProducts();
     }

+    public function it_adds_product_counts_to_anchored_categories_in_collection()
+    {
+        $cat_collection = new \Mage_Catalog_Model_Resource_Category_Collection();
+        // Having to force Magento to believe the collection is loaded to avoid 'item with same id already exists' exception
+        $refObject = new \ReflectionObject($cat_collection);
+        $refCollectionLoaded = $refObject->getProperty('_isCollectionLoaded');
+        $refCollectionLoaded->setAccessible(true);
+        $refCollectionLoaded->setValue($cat_collection, true);
+
+        $category1 = new \Mage_Catalog_Model_Category;
+        $category2 = new \Mage_Catalog_Model_Category;
+        $category1->setData(array('entity_id' => 1, 'is_anchor' => true));
+        $category2->setData(array('entity_id' => 2, 'is_anchor' => true));
+        $cat_collection->addItem($category1)->addItem($category2);
+
+        $params1 = array(
+            'categories' => 1,
+        );
+        $this->_elasticsearch->getProductCount($params1)->willReturn(3);
+
+        $params2 = array(
+            'categories' => 2,
+        );
+        $this->_elasticsearch->getProductCount($params2)->willReturn(4);
+
+        $this->addCountToCategories($cat_collection);
+
+        // $cat_collection->addItem has overkill class introspection preventing the usage of a double
+        // manually having to assert values returned, potentially could use reflection above to force the doubles in.
+        if($category1->getProductCount() != 3){
+            throw new Exception("Value {$category1->getProductCount()} not equal to 3");
+        }
+        if($category2->getProductCount() != 4){
+            throw new Exception("Value {$category2->getProductCount()} not equal to 4");
+        }
+    }
+
 }
diff --git a/src/app/code/local/Magehack/Elasticmage/Model/Resource/Product/Collection.php b/src/app/code/local/Magehack/Elasticmage/Model/Resource/Product/Collection.php
index e211f0f..2e44d3d 100644
--- a/src/app/code/local/Magehack/Elasticmage/Model/Resource/Product/Collection.php
+++ b/src/app/code/local/Magehack/Elasticmage/Model/Resource/Product/Collection.php
@@ -245,8 +245,41 @@ class Magehack_Elasticmage_Model_Resource_Product_Collection extends Mage_Catalo
         return new Varien_Db_Select($this->_getReadAdapter());
     }

-    public function addCountToCategories($argument1)
+    public function addCountToCategories($categoryCollection)
     {
+        $isAnchor    = array();
+        $isNotAnchor = array();
+        foreach ($categoryCollection as $category) {
+            if ($category->getIsAnchor()) {
+                $isAnchor[]    = $category->getId();
+            } else {
+                $isNotAnchor[] = $category->getId();
+            }
+        }
+        $productCounts = array();
+
+        Mage::dispatchEvent(
+            'catalog_product_collection_before_add_count_to_categories',
+            array('collection' => $this)
+        );
+
+        if ($isAnchor) {
+            foreach($isAnchor as $categoryId){
+                $productCounts[$categoryId] = $this->_elasticsearch->getProductCount(array('categories' => $categoryId));
+            }
+        }
+        if ($isNotAnchor) {
+            //TODO: implement isNotAnchor category count
+        }
+
+        foreach ($categoryCollection as $category) {
+            $_count = 0;
+            if (isset($productCounts[$category->getId()])) {
+                $_count = $productCounts[$category->getId()];
+            }
+            $category->setProductCount($_count);
+        }
+
         return $this;
     }