package collection.mutable;

import java.util.AbstractMap;
import java.util.ArrayList;
import java.util.List;
import java.util.HashMap;
import java.util.Map;

public class HashMultimap<K,V> {
    Map<K, List<V>> map = new HashMap<K, List<V>>();
    int total = 0;

    public HashMultimap() { }

    /**
     * Returns the total number of entries in the multimap.
     * @return
     */
    public int size() { return total; }

    /**
     * Tests whether if the multimap is empty.
     * @return
     */
    public boolean isEmpty() { return size() == 0; }

    /**
     * Returns a (possibly empty) iteration of all values associated wiht the key.
     * @param key
     * @return
     */
    Iterable<V> get(K key) {
        List<V> secondary = map.get(key);
        if (secondary != null)
            return secondary;
        return new ArrayList<V>();
    }

    /**
     * Adds a new entry associating key with value.
     * @param key
     * @param value
     */
    public void put(K key, V value) {
        List<V> secondary = map.get(key);
        if (secondary == null) {
            secondary = new ArrayList<V>();
            map.put(key, secondary);
        }
        secondary.add(value);
        total++;
    }

    /**
     * Removes the (key, value) entry, if exists.
     * @param key
     * @param value
     * @return
     */
    public boolean remove(K key, V value) {
        boolean wasRemoved = false;
        List<V> secondary = map.get(key);
        if (secondary != null) {
            wasRemoved = secondary.remove(value);
            if (wasRemoved) {
                total--;
                if (secondary.isEmpty())
                    map.remove(key);
            }
        }
        return wasRemoved;
    }

    /**
     * Removes all entries with given key.
     * @param key
     * @return
     */
    Iterable<V> removeAll(K key) {
        List<V> secondary = map.get(key);
        if (secondary != null) {
            total -= secondary.size();
            map.remove(key);
        } else
            secondary = new ArrayList<V>();
        return secondary;
    }

    /**
     * Returns an iteration of all entries in the multimap.
     * @return
     */
    Iterable<Map.Entry<K,V>> entries() {
        List<Map.Entry<K,V>> result = new ArrayList<Map.Entry<K, V>>();
        for (Map.Entry<K,List<V>> secondary: map.entrySet()) {
            K key = secondary.getKey();
            for (V value: secondary.getValue())
                result.add(new AbstractMap.SimpleEntry<K, V>(key, value));
        }
        return result;
    }
}
