package com.hubspot.mesos;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import com.google.common.collect.Maps;
import com.google.common.primitives.Longs;

public class CounterMap<K> {

  private final Map<K, Counter> map;

  public CounterMap() {
    this.map = new HashMap<K, Counter>();
  }

  public CounterMap(int initialCapacity) {
    this.map = new HashMap<K, Counter>(initialCapacity);
  }

  public Map<K, Long> toCountMap() {
    Map<K, Long> result = Maps.newHashMapWithExpectedSize(map.size());

    for (Entry<K, Counter> entry : map.entrySet()) {
      result.put(entry.getKey(), entry.getValue().count);
    }

    return result;
  }

  public static final class Counter {
    private long count;

    private Counter() {
      count = 0;
    }

    public long getCount() {
      return count;
    }

    @Override
    public String toString() {
      return "Counter [count=" + count + "]";
    }
  }

  public void incr(K key, long amount) {
    Counter c = map.putIfAbsent(key, new Counter());

    c.count += amount;
  }

  public void incr(K key) {
    incr(key, 1);
  }

  public void decr(K key) {
    Counter c = map.putIfAbsent(key, new Counter());

    c.count--;
  }

  public long getCount(K key) {
    Counter c = map.get(key);

    if (c == null) {
      return 0;
    }

    return c.count;
  }

  public Collection<K> getKeys() {
    return map.keySet();
  }

  public void clear() {
    map.clear();
  }

  public List<Entry<K, Counter>> asSortedEntryList() {
    List<Entry<K, Counter>> entries = new ArrayList<Entry<K, Counter>>(map.size());

    for (Entry<K, Counter> entry : map.entrySet()) {
      entries.add(entry);
    }

    Collections.sort(entries, new Comparator<Entry<K, Counter>>() {

      @Override
      public int compare(Entry<K, Counter> o1, Entry<K, Counter> o2) {
        return Longs.compare(o2.getValue().getCount(), o1.getValue().getCount());
      }

    });

    return entries;
  }

  @Override
  public String toString() {
    return "CounterMap{" +
        "map=" + map +
        '}';
  }
}

@@ -443,23 +443,13 @@
 Map%3C
-K, Counter
 %3E();%0A
+
   %7D%0A
@@ -519,26 +519,16 @@
 HashMap%3C
-K, Counter
 %3E(initia
@@ -1117,35 +1117,39 @@
 Counter c = map.
+com
 put
+e
 IfAbsent(key, ne
@@ -1137,32 +1137,37 @@
 uteIfAbsent(key,
+ k -%3E
  new Counter());
@@ -1295,19 +1295,23 @@
 c = map.
+com
 put
+e
 IfAbsent
@@ -1315,16 +1315,21 @@
 ent(key,
+ k -%3E
  new Cou
@@ -1652,24 +1652,24 @@
 tryList() %7B%0A
+
     List%3CEnt
@@ -1708,33 +1708,16 @@
 rayList%3C
-Entry%3CK, Counter%3E
 %3E(map.si

switch to computeIfAbsent to return new key