package macrobase.analysis.summary.count;

import com.google.common.collect.Ordering;
import com.google.common.collect.Sets;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/*
 Maintains probabilistic heavy-hitters:
    - Item counts are overreported.
    - Once we have seen 1/threshold items, size is >= 1/threshold items.

 This is SpaceSaving but with:
   - O(1) update and O(items*log(k)) maintenance
     (normally O(log(k)) update)
   - unlimited space overhead within an epoch (normally O(k))

 Basic idea:
    Store min count from previous epoch
    Item in table? Increment
    Item not in table? Min + count

    End of epoch:
     1.) decay counts
     2.) record new min
     3.) compute 1/thresh highest counts
     4.) discard lower items, updating min if necessary

 */
public class FastButBigSpaceSaving extends ApproximateCount {
    private static final Logger log = LoggerFactory.getLogger(FastButBigSpaceSaving.class);

    private HashMap<Integer, Double> counts = new HashMap<>();
    private double totalCount = 0;
    private final double threshold;
    private final int maxStableSize;

    private double prevEpochMin = 0;

    // threshold == totalCount*threshold;
    public FastButBigSpaceSaving(double threshold) {
        this.threshold = threshold;
        maxStableSize = (int)(1/threshold);
    }

    @Override
    public void multiplyAllCounts(Double by) {
        totalCount *= by;

        log.trace("Decaying; {} items stored", counts.size());

        for(Map.Entry<Integer, Double> entry : counts.entrySet()) {
            double newValue = entry.getValue()*by;
            counts.put(entry.getKey(), newValue);

        }

        if(counts.size() > maxStableSize) {
            Set<Map.Entry<Integer, Double>> frozenEntrySet = Sets.newHashSet(counts.entrySet());

            // find the <maxStableSize>th value
            double threshValue = Ordering.natural().greatestOf(counts.values(), maxStableSize).get(maxStableSize-1);
            prevEpochMin = threshValue;

            log.trace("Pruning! Minimum threshold for count is {} ({})", threshValue, threshold);

            for(Map.Entry<Integer, Double> entry : frozenEntrySet) {
                if(entry.getValue() < threshValue) {
                    counts.remove(entry.getKey());
                }
            }
        }

        log.trace("Finished pruning; new size is {}", counts.size());
    }

    public HashMap<Integer, Double> getCounts() {
        return counts;
    }

    @Override
    public void observe(Integer item, double count) {
        Double value = counts.get(item);
        if (value == null) {
            value = prevEpochMin + count;
            totalCount += prevEpochMin + count;
        } else {
           value += count;
           totalCount += count;
        }

        counts.put(item, value);
    }

    @Override
    public double getTotalCount() {
        return totalCount;
    }

    @Override
    public double getCount(int item) {
        Double ret = counts.get(item);
        if(ret == null) {
            return 0;
        }

        return ret;
    }
}

@@ -36,16 +36,110 @@
 count;%0A%0A
+import com.google.common.collect.Lists;%0Aimport com.google.common.collect.MinMaxPriorityQueue;%0A
 import c
@@ -291,37 +291,121 @@
 til.
-HashMap;%0Aimport java.util.Map
+Comparator;%0Aimport java.util.HashMap;%0Aimport java.util.List;%0Aimport java.util.Map;%0Aimport java.util.PriorityQueue
 ;%0Aim
@@ -1974,18 +1974,19 @@
         
-Se
+Lis
 t%3CMap.En
@@ -2011,39 +2011,29 @@
 e%3E%3E 
-frozenEntrySet = Sets.newHashSe
+a = Lists.newArrayLis
 t(co
@@ -2046,25 +2046,24 @@
 ntrySet());%0A
-%0A
             
@@ -2066,43 +2066,68 @@
     
-// find the %3CmaxStableSize%3Eth value
+a.sort((e1, e2) -%3E e1.getValue().compareTo(e2.getValue()));%0A
 %0A   
@@ -2146,88 +2146,65 @@
 ble 
-thresh
+prev
 Val
-ue
  = 
-Ordering.natural().greatestOf(counts.values(), maxStableSize).get(
+-1;%0A            int toRemove = counts.size() - 
 maxS
@@ -2216,12 +2216,10 @@
 Size
--1)
 ;
+%0A
 %0A   
@@ -2231,136 +2231,112 @@
     
-prevEpochMin = threshValue;%0A%0A            log.trace(%22Pruning! Minimum threshold for count is %7B%7D (%7B%7D)%22, threshValue, threshold);%0A%0A
+log.trace(%22Removing %7B%7D items from counts%22, toRemove);%0A%0A            for(int i = 0; i %3C toRemove; ++i) %7B%0A 
     
@@ -2335,36 +2335,35 @@
  %7B%0A             
-for(
+   
 Map.Entry%3CIntege
@@ -2383,115 +2383,101 @@
 try 
-: frozenEntrySet) %7B%0A                if(entry.getValue() %3C threshValue) %7B%0A                    counts.remove(
+= a.get(i);%0A                counts.remove(entry.getKey());%0A                assert (prevVal %3C 
 entr
@@ -2473,35 +2473,37 @@
 vVal %3C entry.get
-Key
+Value
 ());%0A           
@@ -2499,33 +2499,59 @@
                 
-%7D
+prevVal = entry.getValue();
 %0A            %7D%0A 

Use our own sort for FastButBigSpaceSaving pruning.