package ru.compscicenter.java2015.collections;

import java.util.AbstractCollection;
import java.util.Collection;
import java.util.ConcurrentModificationException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

/**
 * Created by qurbonzoda on 07.11.15.
 */
public class SimpleMultiSet<E> extends AbstractCollection<E> implements MultiSet<E> {
    private Map<E, Integer> elements;
    private int size;
    public SimpleMultiSet() {
        elements = new HashMap<>();
    }
    public SimpleMultiSet(Collection<? extends E> collection) {
        elements = new HashMap<>();
        for (E e : collection) {
            add(e);
        }
    }
    @Override
    public int size() {
        return size;
    }
    @Override
    public Iterator<E> iterator() {
        return new SimpleMultiSetItetator();
    }
    @Override
    public boolean add(E e) {
        add(e, 1);
        return true;
    }
    @Override
    public int add(E e, int occurrences) {
        if (occurrences < 0) {
            throw new IllegalArgumentException();
        }
        if (elements.containsKey(e)) {
            elements.put(e, elements.get(e) + occurrences);
        } else {
            elements.put(e, occurrences);
        }
        size += occurrences;
        return count(e) - occurrences;
    }
    @Override
    public int remove(Object o, int occurences) {
        if (occurences < 0) {
            throw new IllegalArgumentException();
        }
        int cnt = count(o);
        if (elements.containsKey(o)) {
            if (occurences >= cnt) {
                size -= cnt;
                elements.remove(o);
            } else {
                size -= occurences;
                elements.put((E) o, cnt - occurences);
            }
        }
        return cnt;
    }
    @Override
    public boolean remove(Object o) {
        return 0 < remove(o, 1);
    }
    @Override
    public int count(Object o) {
        if (elements.containsKey(o)) {
            return elements.get(o);
        }
        return 0;
    }
    @Override
    public boolean contains(Object o) {
        return elements.containsKey(o);
    }
    @Override
    public boolean removeAll(Collection<?> c) {
        boolean anyRemoved = false;
        for (Object o : c) {
            anyRemoved = remove(o, Integer.MAX_VALUE) > 0;
        }
        return anyRemoved;
    }
    @Override
    public boolean retainAll(Collection<?> c) {
        boolean anyRemoved = false;
        Iterator<Map.Entry<E, Integer>> iter = elements.entrySet().iterator();
        while (iter.hasNext()) {
            Map.Entry<E, Integer> entry = iter.next();
            if (!c.contains(entry.getKey())) {
                size -= entry.getValue();
                anyRemoved = true;
                iter.remove();
            }
        }
        return anyRemoved;
    }
    @Override
    public void clear() {
        size = 0;
        elements.clear();
    }
    @Override
    public boolean equals(Object o) {
        if (o != null && o instanceof SimpleMultiSet) {
            SimpleMultiSet temp = (SimpleMultiSet) o;
            if (temp.size() != size()) {
                return false;
            }
            for (Map.Entry<E, Integer> entry : elements.entrySet()) {
                if (temp.count(entry.getKey()) != entry.getValue()) {
                    return false;
                }
            }
            return true;
        }
        return false;
    }

    class SimpleMultiSetItetator implements Iterator<E> {
        private final Iterator<Map.Entry<E, Integer>> entryIterator;
        private int occurenceLeft;
        private boolean canRemove;
        private Map.Entry<E, Integer> currentEntry;
        SimpleMultiSetItetator() {
            entryIterator = elements.entrySet().iterator();
            occurenceLeft = 0;
        }
        @Override
        public boolean hasNext() {
            return occurenceLeft > 0 || entryIterator.hasNext();
        }
        @Override
        public E next() {
            if (occurenceLeft == 0) {
                currentEntry = entryIterator.next();
                occurenceLeft = currentEntry.getValue();
            }
            occurenceLeft--;
            canRemove = true;
            return currentEntry.getKey();
        }
        @Override
        public void remove() {
            if (!canRemove) {
                throw new ConcurrentModificationException();
            }
            if (currentEntry.getValue() == 1) {
                entryIterator.remove();
                size--;
            } else {
                SimpleMultiSet.this.remove(currentEntry.getKey());
            }
            canRemove = false;
        }
    }
}