package ru.yandex.webmaster.common.util;

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

import org.apache.commons.lang.Validate;

/**
 * Implementation of "Heavy Hitters: Stream-Summary"
 * A. Metwally, D. Agrawal, A.E. Abbadi. Efficient Computation of Frequent and Top-K Elements in Data Streams
 *
 * @author aherman
 */
public class TopKCounter<T> {
    private final int maxMeasures;

    private final Bucket<T> head;
    private final Bucket<T> tail;

    private final Map<T, InternalMeasure<T>> entryToMeasure = new HashMap<>();

    public TopKCounter(int maxMeasures) {
        this.maxMeasures = maxMeasures;

        head = new Bucket<>(0);
        tail = new Bucket<>(Integer.MAX_VALUE);

        head.next = tail;
        tail.previous = head;
    }

    public Measure<T> getMeasure(T entry) {
        InternalMeasure<T> internalMeasure = entryToMeasure.get(entry);
        if (internalMeasure != null) {
            return new Measure<>(entry, internalMeasure.bucket.count, internalMeasure.error);
        }

        Bucket<T> smallestBucket = head.next;
        if (smallestBucket == tail) {
            return new Measure<>(entry, 0, 0);
        }
        return new Measure<>(entry, smallestBucket.count, smallestBucket.count);
    }

    public Measure<T> newEntry(T entry) {
        InternalMeasure<T> measure = entryToMeasure.get(entry);
        Bucket<T> bucket = null;
        int count;
        if (measure != null) {
            bucket = measure.getBucket();
            count = bucket.count + 1;
            bucket.moveMeasure(measure, count);
        } else {
            if (entryToMeasure.size() < maxMeasures) {
                measure = new InternalMeasure<>(entry, 0, head);
                count = 1;
                head.moveMeasure(measure, count);
            } else {
                InternalMeasure<T> removedMeasure = head.next.removeFirst();
                entryToMeasure.remove(removedMeasure.entry);

                count = head.next.count;
                measure = new InternalMeasure<>(entry, count, head);
                count += 1;
                head.next.moveMeasure(measure, count);
                bucket = head.next;
            }
            entryToMeasure.put(entry, measure);
        }
        if (bucket != null && bucket.isEmpty()) {
            bucket.remove();
        }
        return new Measure<>(entry, count, measure.error);
    }

    public Iterator<Measure<T>> topItemIterator() {
        Bucket<T> currentBucket = tail;
        while (currentBucket != head && currentBucket.measures.isEmpty()) {
            currentBucket = currentBucket.previous;
        }
        if (currentBucket == head) {
            return Collections.emptyIterator();
        }

        return new TopItemIterator<>(currentBucket);
    }

    public static class Measure<T> {
        private final T entry;
        private final int count;
        private final int error;

        public Measure(T entry, int count, int error) {
            this.entry = entry;
            this.count = count;
            this.error = error;
        }

        public int getCount() {
            return count;
        }

        public int getError() {
            return error;
        }

        public T getEntry() {
            return entry;
        }
    }

    private static class InternalMeasure<T> {
        private final T entry;
        private final int error;
        private Bucket<T> bucket;

        private InternalMeasure(T entry, int error, Bucket<T> bucket) {
            this.entry = entry;
            this.error = error;
            this.bucket = bucket;
        }

        public int getError() {
            return error;
        }

        public Bucket<T> getBucket() {
            return bucket;
        }

        public void setBucket(Bucket<T> bucket) {
            this.bucket = bucket;
        }

        @Override
        public String toString() {
            return "Measure " + bucket.count + "/" + error;
        }
    }

    private static class Bucket<T> {
        private final int count;
        private final Set<InternalMeasure<T>> measures = new HashSet<>();

        private Bucket<T> next;
        private Bucket<T> previous;

        private Bucket(int count) {
            this.count = count;
        }

        public void insertNext(Bucket<T> bucket) {
            Validate.isTrue(bucket.count > count);
            Bucket<T> oldNext = this.next;

            bucket.next = oldNext;
            this.next = bucket;

            oldNext.previous = bucket;
            bucket.previous = this;
        }

        public void remove() {
            previous.next = next;
            next.previous = previous;
        }

        public void moveMeasure(InternalMeasure<T> m, int measureCount) {
            measures.remove(m);
            if (measureCount == next.count) {
                next.addEntry(m);
            } else {
                Bucket<T> b = new Bucket<>(measureCount);
                insertNext(b);
                b.addEntry(m);
            }
        }

        public void addEntry(InternalMeasure<T> entry) {
            measures.add(entry);
            entry.setBucket(this);
        }

        public boolean isEmpty() {
            return measures.isEmpty();
        }

        public InternalMeasure<T> removeFirst() {
            InternalMeasure<T> m = measures.iterator().next();
            measures.remove(m);
            return m;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }

            Bucket bucket = (Bucket) o;

            if (count != bucket.count) {
                return false;
            }

            return true;
        }

        @Override
        public int hashCode() {
            return count;
        }

        @Override
        public String toString() {
            return "Bucket [" + count + "]: " + measures.size();
        }
    }

    private class TopItemIterator<T> implements Iterator<Measure<T>> {
        private Bucket<T> currentBucket;
        private Iterator<InternalMeasure<T>> it;
        private int selfCheckCounter = 0;

        public TopItemIterator(Bucket<T> currentBucket) {
            this.currentBucket = currentBucket;
            it = currentBucket.measures.iterator();
        }

        @Override
        public boolean hasNext() {
            return it.hasNext();
        }

        @Override
        public Measure<T> next() {
            InternalMeasure<T> internalMeasure = it.next();
            Measure<T> result = new Measure<>(internalMeasure.entry, internalMeasure.bucket.count, internalMeasure.error);
            while (!it.hasNext()) {
                currentBucket = currentBucket.previous;
                if (currentBucket == null) {
                    it = Collections.emptyIterator();
                    break;
                } else {
                    it = currentBucket.measures.iterator();
                }
                // Should never happen
                if (selfCheckCounter++ > maxMeasures) {
                    throw new IllegalStateException("Infinite cycle possibility: too many buckets");
                }
            }
            return result;
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException();
        }
    }
}
