#pragma once
#include <BWAPI.h>

class ArmyComposition
{
public:
  void addUnlimited(BWAPI::UnitType unitType, size_t nb)
  {
    m_unlimitedComposition[unitType] += nb;
  }

  void addLimited(BWAPI::UnitType unitType, size_t nb)
  {
    m_limitedComposition[unitType] += nb;
  }

  //typedef std::map<BWAPI::UnitType, size_t>::const_iterator const_iterator;
  //const_iterator begin() const { return m_unlimitedComposition.begin(); }
  //const_iterator end() const { return m_unlimitedComposition.end(); }

  size_t getNbUnlimited(BWAPI::UnitType unitType) const
  {
    std::map<BWAPI::UnitType, size_t>::const_iterator it = m_unlimitedComposition.find(unitType);
    if (it == m_unlimitedComposition.end())
      return 0;
    else
      return it->second;
  }
  size_t getNbLimited(BWAPI::UnitType unitType) const
  {
    std::map<BWAPI::UnitType, size_t>::const_iterator it = m_limitedComposition.find(unitType);
    if (it == m_limitedComposition.end())
      return 0;
    else
      return it->second;
  }

  bool operator==(const ArmyComposition& other) const
  {
    return m_unlimitedComposition == other.m_unlimitedComposition &&
           m_limitedComposition == other.m_limitedComposition;
  }
  bool operator!=(const ArmyComposition& other) const
  {
    return m_unlimitedComposition != other.m_unlimitedComposition ||
           m_limitedComposition != other.m_limitedComposition;
  }
  const std::map<BWAPI::UnitType, size_t>& getLimitedComposition() const { return m_limitedComposition; }
  const std::map<BWAPI::UnitType, size_t>& getUnlimitedComposition() const { return m_unlimitedComposition; }

protected:
  std::map<BWAPI::UnitType, size_t> m_limitedComposition;
  std::map<BWAPI::UnitType, size_t> m_unlimitedComposition;
};