CounterMap.java

package edu.berkeley.nlp.util;


import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

//import fig.basic.Pair;

/**
 * Maintains counts of (key, value) pairs.  The map is structured so that for
 * every key, one can get a counter over values.  Example usage: keys might be
 * words with values being POS tags, and the count being the number of
 * occurences of that word/tag pair.  The sub-counters returned by
 * getCounter(word) would be count distributions over tags for that word.
 *
 * @author Dan Klein
 */
public class CounterMap<K, V> implements java.io.Serializable {
	private static final long serialVersionUID = 1L;
	MapFactory<V, Double> mf;
	Map<K, Counter<V>> counterMap;
    double defltVal = 0.0;

	protected Counter<V> ensureCounter(K key) {
		Counter<V> valueCounter = counterMap.get(key);
		if (valueCounter == null) {
			valueCounter = buildCounter(mf);
            valueCounter.setDeflt(defltVal);
			counterMap.put(key, valueCounter);
		}
		return valueCounter;
	}

  public Collection<Counter<V>> getCounters() {
    return counterMap.values();
  }

	/**
	 * @return
	 */
	protected Counter<V> buildCounter(MapFactory<V, Double> mf)
	{
		return new Counter<V>(mf);
	}

	/**
	 * Returns the keys that have been inserted into this CounterMap.
	 */
	public Set<K> keySet() {
		return counterMap.keySet();
	}

	/**
	 * Sets the count for a particular (key, value) pair.
	 */
	public void setCount(K key, V value, double count) {
		Counter<V> valueCounter = ensureCounter(key);
		valueCounter.setCount(value, count);
	}
	
//	public void setCount(Pair<K,V> pair) {
//		
//	}

	/**
	 * Increments the count for a particular (key, value) pair.
	 */
	public void incrementCount(K key, V value, double count) {
		Counter<V> valueCounter = ensureCounter(key);
		valueCounter.incrementCount(value, count);
	}

	/**
	 * Gets the count of the given (key, value) entry, or zero if that entry is
	 * not present.  Does not create any objects.
	 */
	public double getCount(K key, V value) {
		Counter<V> valueCounter = counterMap.get(key);
		if (valueCounter == null) return defltVal;
		return valueCounter.getCount(value);
	}

	/**
	 * Gets the sub-counter for the given key.  If there is none, a counter is
	 * created for that key, and installed in the CounterMap.  You can, for
	 * example, add to the returned empty counter directly (though you shouldn't).
	 * This is so whether the key is present or not, modifying the returned
	 * counter has the same effect (but don't do it).
	 */
	public Counter<V> getCounter(K key) {
		return ensureCounter(key);
	}

	public void incrementAll(Map<K, V> map, double count) {
		for (Map.Entry<K, V> entry : map.entrySet()) {
			incrementCount(entry.getKey(), entry.getValue(), count);
		}
	}

	public void incrementAll(CounterMap<K, V> cMap) {
		for (Map.Entry<K,Counter<V>> entry: cMap.counterMap.entrySet()) {
      K key = entry.getKey();
      Counter<V> innerCounter = entry.getValue();
      for (Map.Entry<V, Double> innerEntry: innerCounter.entrySet()) {
        V value = innerEntry.getKey();
        incrementCount(key,value,innerEntry.getValue());
      }
		}
	}

	/**
	 * Gets the total count of the given key, or zero if that key is
	 * not present.  Does not create any objects.
	 */
	public double getCount(K key) {
		Counter<V> valueCounter = counterMap.get(key);
		if (valueCounter == null) return 0.0;
		return valueCounter.totalCount();
	}

	/**
	 * Returns the total of all counts in sub-counters.  This implementation is
	 * linear; it recalculates the total each time.
	 */
	public double totalCount() {
		double total = 0.0;
		for (Map.Entry<K, Counter<V>> entry : counterMap.entrySet()) {
			Counter<V> counter = entry.getValue();
			total += counter.totalCount();
		}
		return total;
	}

	/**
	 * Returns the total number of (key, value) entries in the CounterMap (not
	 * their total counts).
	 */
	public int totalSize() {
		int total = 0;
		for (Map.Entry<K, Counter<V>> entry : counterMap.entrySet()) {
			Counter<V> counter = entry.getValue();
			total += counter.size();
		}
		return total;
	}

	/**
	 * The number of keys in this CounterMap (not the number of key-value entries
	 * -- use totalSize() for that)
	 */
	public int size() {
		return counterMap.size();
	}

	/**
	 * True if there are no entries in the CounterMap (false does not mean
	 * totalCount > 0)
	 */
	public boolean isEmpty() {
		return size() == 0;
	}

	/**
	 * Finds the key with maximum count.  This is a linear operation, and ties are broken arbitrarily.
	 *
	 * @return a key with minumum count
	 */
	public Pair<K, V> argMax() {
		double maxCount = Double.NEGATIVE_INFINITY;
		Pair<K, V> maxKey = null;
		for (Map.Entry<K, Counter<V>> entry : counterMap.entrySet()) {
			Counter<V> counter = entry.getValue();
			V localMax = counter.argMax();
			if (counter.getCount(localMax) > maxCount || maxKey == null) {
				maxKey = new Pair<K, V>(entry.getKey(), localMax);
				maxCount = counter.getCount(localMax);
			}
		}
		return maxKey;
	}


	public String toString(int maxValsPerKey) {
		StringBuilder sb = new StringBuilder("[\n");
		for (Map.Entry<K, Counter<V>> entry : counterMap.entrySet()) {
			sb.append("  ");
			sb.append(entry.getKey());
			sb.append(" -> ");
			sb.append(entry.getValue().toString(maxValsPerKey));
			sb.append("\n");
		}
		sb.append("]");
		return sb.toString();
	}

  @Override
  public String toString() {
    return toString(20);
  }

  public String toString(Collection<String> keyFilter) {
    StringBuilder sb = new StringBuilder("[\n");
    for (Map.Entry<K, Counter<V>> entry : counterMap.entrySet()) {
      if (keyFilter != null && !keyFilter.contains(entry.getKey())) {
        continue;
      }
      sb.append("  ");
      sb.append(entry.getKey());
      sb.append(" -> ");
      sb.append(entry.getValue().toString(20));
      sb.append("\n");
    }
    sb.append("]");
    return sb.toString();
  }


	public CounterMap(CounterMap<K,V> cm)
	{
		this();
		incrementAll(cm);
	}
	
	public CounterMap() {
		this(false);
	}
	
	
	public boolean isEqualTo(CounterMap<K,V> map)
	{
		boolean tmp = true;
		CounterMap<K,V> bigger = map.size() > size() ? map : this;
		for (K k : bigger.keySet())
		{
			tmp &= map.getCounter(k).isEqualTo(getCounter(k));
		}
		return tmp;
	}
	
	
	
	

	public CounterMap(MapFactory<K, Counter<V>> outerMF, MapFactory<V, Double> innerMF) {
		mf = innerMF;
		counterMap = outerMF.buildMap();
	}

	public CounterMap(boolean identityHashMap) {
		this(identityHashMap ? new MapFactory.IdentityHashMapFactory<K, Counter<V>>()
				: new MapFactory.HashMapFactory<K, Counter<V>>(),
				identityHashMap ? new MapFactory.IdentityHashMapFactory<V, Double>()
						: new MapFactory.HashMapFactory<V, Double>());
	}

	public static void main(String[] args) {
		CounterMap<String, String> bigramCounterMap = new CounterMap<String, String>();
		bigramCounterMap.incrementCount("people", "run", 1);
		bigramCounterMap.incrementCount("cats", "growl", 2);
		bigramCounterMap.incrementCount("cats", "scamper", 3);
		System.out.println(bigramCounterMap);
		System.out.println("Entries for cats: " + bigramCounterMap.getCounter("cats"));
		System.out.println("Entries for dogs: " + bigramCounterMap.getCounter("dogs"));
		System.out.println("Count of cats scamper: "
				+ bigramCounterMap.getCount("cats", "scamper"));
		System.out.println("Count of snakes slither: "
				+ bigramCounterMap.getCount("snakes", "slither"));
		System.out.println("Total size: " + bigramCounterMap.totalSize());
		System.out.println("Total count: " + bigramCounterMap.totalCount());
		System.out.println(bigramCounterMap);
	}

	public void normalize() {
		for (K key : keySet()) {
			getCounter(key).normalize();
		}
	}
	
	public void normalizeWithDiscount(double discount) {
		for (K key : keySet()) {
			Counter<V> ctr = getCounter(key);
			double totalCount = ctr.totalCount();
			for (V value : ctr.keySet()) {
				ctr.setCount(value, (ctr.getCount(value) - discount) / totalCount);
			}
		}
	}
	
	/**
	 * Constructs reverse CounterMap where the count of a pair (k,v)
	 * is the count of (v,k) in the current CounterMap
	 * @return
	 */
	public CounterMap<V,K> invert() {
		CounterMap<V,K> invertCounterMap = new CounterMap<V, K>();
		for (K key: this.keySet()) {
			Counter<V> keyCounts = this.getCounter(key);
			for (V val: keyCounts.keySet()) {
				double count = keyCounts.getCount(val);
				invertCounterMap.setCount(val, key, count);
			}
		}
		return invertCounterMap;
	}

	/**
	 * Scale all entries in <code>CounterMap</code>
	 * by <code>scaleFactor</code>
	 * @param scaleFactor
	 */
	public void scale(double scaleFactor) {
		for (K key: keySet()) {
		  Counter<V> counts = getCounter(key);
			counts.scale(scaleFactor);
		}
	}

	public boolean containsKey(K key) {
		return counterMap.containsKey(key);
	}
	
	public Iterator<Pair<K,V>> getPairIterator() {
		
		class PairIterator implements Iterator<Pair<K,V>> {
			
			Iterator<K> outerIt ;
			Iterator<V> innerIt ;
			K curKey ;
			
			public PairIterator() {
				outerIt = keySet().iterator();
			}
			
			private boolean advance() {				
				if (innerIt == null || !innerIt.hasNext()) {
					if (!outerIt.hasNext()) {
						return false;
					}
					curKey = outerIt.next();
					innerIt = getCounter(curKey).keySet().iterator();
				} 	
				return true;
			}

			public boolean hasNext() {
				return advance();
			}

			public Pair<K, V> next() {
				advance();
				assert curKey != null;
				return Pair.newPair(curKey, innerIt.next());
			}

			public void remove() {
				// TODO Auto-generated method stub
				
			}
			
		};
		return new PairIterator();
	}

	public Set<Map.Entry<K, Counter<V>>> getEntrySet() {
		// TODO Auto-generated method stub
		return counterMap.entrySet();
	}

	public void removeKey(K oldIndex)
	{
		counterMap.remove(oldIndex);		
	}

	public void setCounter(K newIndex, Counter<V> counter)
	{
		counterMap.put(newIndex, counter);
		
	}

	public void setDefault(double defltVal) {
    this.defltVal = defltVal;
    for (Counter<V> vCounter : counterMap.values()) {
      vCounter.setDeflt(defltVal);
    }
  }
	
}