/*
 * Copyright (C) 2007 Google Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.common.collect;

import com.google.common.base.Nullable;
import com.google.common.base.Objects;

import static com.google.common.base.Preconditions.checkArgument;

import java.io.Serializable;
import java.util.AbstractCollection;
import java.util.AbstractMap;
import java.util.AbstractSet;
import java.util.Collection;
import java.util.Comparator;
import java.util.ConcurrentModificationException;
import java.util.Iterator;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.RandomAccess;
import java.util.Set;
import java.util.SortedMap;
import java.util.SortedSet;

/**
 * Basic implementation of the {@link Multimap} interface. This class represents
 * a multimap as a map that associates each key with a collection of values. All
 * methods of {@link Multimap} are supported, including those specified as
 * optional in the interface.
 *
 * <p>To implement a multimap, a subclass must define the method {@link
 * #createCollection()}, which creates an empty collection of values for a key.
 *
 * <p>The multimap constructor takes a map that has a single entry for each
 * distinct key. When you insert a key-value pair with a key that isn't already
 * in the multimap, {@code StandardMultimap} calls {@link #createCollection()}
 * to create the collection of values for that key. The subclass should not call
 * {@link #createCollection()} directly, and a new instance should be created
 * every time the method is called.
 *
 * <p>For example, the subclass could pass a {@link java.util.TreeMap} during
 * construction, and {@link #createCollection()} could return a {@link
 * java.util.TreeSet}, in which case the multimap's iterators would propagate
 * through the keys and values in sorted order.
 *
 * <p>Keys and values may be null, as long as the underlying collection classes
 * support null elements.
 *
 * <p>The collections created by {@link #createCollection()} may or may not
 * allow duplicates. If the collection, such as a {@link Set}, does not support
 * duplicates, an added key-value pair will replace an existing pair with the
 * same key and value, if such a pair is present. With collections like {@link
 * List} that allow duplicates, the collection will keep the existing key-value
 * pairs while adding a new pair.
 *
 * <p>This class is not threadsafe when any concurrent operations update the
 * multimap, even if the underlying map and {@link #createCollection()} method
 * return threadsafe classes. Concurrent read operations will work correctly. To
 * allow concurrent update operations, wrap your multimap with a call to {@link
 * Multimaps#synchronizedMultimap}.
 *
 * @author Jared Levy
 */
abstract class StandardMultimap<K, V> implements Multimap<K, V>, Serializable {
  /*
   * Here's an outline of the overall design.
   *
   * The map variable contains the collection of values associated with each
   * key. When a key-value pair is added to a multimap that didn't previously
   * contain and values for that key, a new collection generated by
   * createCollection is added to the map. That same collection instance
   * remains in the map as long as the multimap has any values for the key. If
   * all values for the key are removed, the key and collection are removed
   * from the map.
   *
   * The get method returns a WrappedCollection, which decorates the collection
   * in the map (if the key is present) or an empty collection (if the key is
   * not present). When the collection delegate in the WrappedCollection is
   * empty, the multimap may contain subsequently added values for that key. To
   * handle that situation, the WrappedCollection checks whether map contains
   * an entry for the provided key, and if so replaces the delegate.
   */

  private final Map<K, Collection<V>> map;
  private int totalSize;

  /**
   * Creates a new multimap that uses the provided map.
   *
   * @param map place to store the mapping from each key to its corresponding
   *     values
   * @throws IllegalArgumentException if {@code map} is not empty
   */
  protected StandardMultimap(Map<K, Collection<V>> map) {
    checkArgument(map.isEmpty());
    this.map = map;
  }

  /**
   * Creates the collection of values for a single key.
   *
   * <p>Collections with weak, soft, or phantom references are not supported.
   * Each call to {@code createCollection} should create a new instance.
   *
   * <p>The returned collection class determines whether duplicate key-value
   * pairs are allowed.
   *
   * @return an empty collection of values
   */
  abstract Collection<V> createCollection();

  /**
   * Creates the collection of values for an explicitly provided key. By
   * default, it simply calls {@link #createCollection()}, which is the correct
   * behavior for most implementations. The {@link LinkedHashMultimap} class
   * overrides it. 
   *
   * @param key key to associate with values in the collection
   * @return an empty collection of values
   */
  Collection<V> createCollection(@Nullable K key) {
    return createCollection();
  }

  Map<K, Collection<V>> backingMap() {
    return map;
  }
  
  // Query Operations
  
  public int size() {
    return totalSize;
  }

  public boolean isEmpty() {
    return totalSize == 0;
  }

  public boolean containsKey(@Nullable Object key) {
    return map.containsKey(key);
  }

  public boolean containsValue(@Nullable Object value) {
    for (Collection<V> collection : map.values()) {
      if (collection.contains(value)) {
        return true;
      }
    }

    return false;
  }

  public boolean containsEntry(@Nullable Object key, @Nullable Object value) {
    Collection<V> collection = map.get(key);
    return collection != null && collection.contains(value);
  }

  // Modification Operations  
  
  public boolean put(@Nullable K key, @Nullable V value) {
    Collection<V> collection = getOrCreateCollection(key);

    if (collection.add(value)) {
      totalSize++;
      return true;
    } else {
      return false;
    }
  }

  private Collection<V> getOrCreateCollection(K key) {
    Collection<V> collection = map.get(key);
    if (collection == null) {
      collection = createCollection(key);
      map.put(key, collection);
    }
    return collection;
  }

  public boolean remove(@Nullable Object key, @Nullable Object value) {
    Collection<V> collection = map.get(key);
    if (collection == null) {
      return false;
    }

    boolean changed = collection.remove(value);
    if (changed) {
      totalSize--;
      if (collection.isEmpty()) {
        map.remove(key);
      }
    }
    return changed;
  }

  // Bulk Operations

  public void putAll(@Nullable K key, Iterable<? extends V> values) {
    if (!values.iterator().hasNext()) {
      return;
    }
    Collection<V> collection = getOrCreateCollection(key);
    int oldSize = collection.size();

    if (values instanceof Collection<?>) {
      @SuppressWarnings("unchecked")
      Collection<? extends V> c = (Collection<? extends V>) values;
      collection.addAll(c);
    } else {
      for (V value : values) {
        collection.add(value);
      }
    }

    totalSize += (collection.size() - oldSize);
  }

  public void putAll(Multimap<? extends K, ? extends V> multimap) {
    for (Map.Entry<? extends K, ? extends V> entry : multimap.entries()) {
      put(entry.getKey(), entry.getValue());
    }
  }

  public Collection<V> replaceValues(
      @Nullable K key, Iterable<? extends V> values) {
    Iterator<? extends V> iterator = values.iterator();
    if (!iterator.hasNext()) {
      return removeAll(key);
    }

    Collection<V> collection = getOrCreateCollection(key);
    Collection<V> oldValues = createCollection();
    oldValues.addAll(collection);

    totalSize -= collection.size();
    collection.clear();

    while (iterator.hasNext()) {
      if (collection.add(iterator.next())) {
        totalSize++;
      }
    }

    return oldValues;
  }

  public Collection<V> removeAll(@Nullable Object key) {
    Collection<V> collection = map.remove(key);
    Collection<V> output = createCollection();

    if (collection != null) {
      output.addAll(collection);
      totalSize -= collection.size();
      collection.clear();
    }

    return output;
  }

  public void clear() {
    // Clear each collection, to make previously returned collections empty.
    for (Collection<V> collection : map.values()) {
      collection.clear();
    }
    map.clear();
    totalSize = 0;
  }

  // Views
  
  /**
   * {@inheritDoc}
   * 
   * <p>The returned collection is not serializable.
   */
  public Collection<V> get(@Nullable K key) {
    Collection<V> collection = map.get(key);
    if (collection == null) {
      collection = createCollection(key);
    }
    return wrapCollection(key, collection);
  }

  /**
   * Generates a decorated collection that remains consistent with the values in
   * the multimap for the provided key. Changes to the multimap may alter the
   * returned collection, and vice versa.
   */
  private Collection<V> wrapCollection(
      @Nullable K key, Collection<V> collection) {
    if (collection instanceof SortedSet<?>) {
      return new WrappedSortedSet(key, (SortedSet<V>) collection, null);
    } else if (collection instanceof Set<?>) {
      return new WrappedSet(key, (Set<V>) collection);
    } else if (collection instanceof List<?>) {
      return wrapList(key, (List<V>) collection, null);
    } else {
      return new WrappedCollection(key, collection, null);
    }
  }

  private List<V> wrapList(
      @Nullable K key, List<V> list, @Nullable WrappedCollection ancestor) {
    return (list instanceof RandomAccess)
        ? new RandomAccessWrappedList(key, list, ancestor)
        : new WrappedList(key, list, ancestor);
  }

  /**
   * Collection decorator that stays in sync with the multimap values for a key.
   * There are two kinds of wrapped collections: full and subcollections. Both
   * have a delegate pointing to the underlying collection class.
   *
   * <p>Full collections, identified by a null ancestor field, contain all
   * multimap values for a given key. Its delegate is a value in {@link
   * StandardMultimap#map} whenever the delegate is non-empty. The {@code
   * refreshIfEmpty}, {@code removeIfEmpty}, and {@code addToMap} methods ensure
   * that the WrappedCollection and map remain consistent.
   *
   * <p>A subcollection, such as a sublist, contains some of the values for a
   * given key. Its ancestor field points to the full wrapped collection with
   * all values for the key. The subcollection {@code refreshIfEmpty}, {@code
   * removeIfEmpty}, and {@code addToMap} methods call the corresponding methods
   * of the full wrapped collection.
   */
  private class WrappedCollection extends AbstractCollection<V> {
    final K key;
    Collection<V> delegate;
    final WrappedCollection ancestor;
    final Collection<V> originalAncestorDelegate;

    WrappedCollection(@Nullable K key, Collection<V> delegate,
        @Nullable WrappedCollection ancestor) {
      this.key = key;
      this.delegate = delegate;
      this.ancestor = ancestor;
      this.originalAncestorDelegate
          = (ancestor == null) ? null : ancestor.getDelegate();
    }

    /**
     * If the delegate collection is empty, but the multimap has values for the
     * key, replace the delegate with the new collection for the key.
     *
     * <p>For a subcollection, refresh its ancestor and validate that the
     * ancestor delegate hasn't changed.
     */
    void refreshIfEmpty() {
      if (ancestor != null) {
        ancestor.refreshIfEmpty();
        if (ancestor.getDelegate() != originalAncestorDelegate) {
          throw new ConcurrentModificationException();
        }
      } else if (delegate.isEmpty()) {
        Collection<V> newDelegate = map.get(key);
        if (newDelegate != null) {
          delegate = newDelegate;
        }
      }
    }

    /**
     * If collection is empty, remove it from {@code map}. For subcollections,
     * check whether the ancestor collection is empty.
     */
    void removeIfEmpty() {
      if (ancestor != null) {
        ancestor.removeIfEmpty();
      } else if (delegate.isEmpty()) {
        map.remove(key);
      }
    }

    K getKey() {
      return key;
    }

    /**
     * Add the delegate to the map. Other {@code WrappedCollection} methods
     * should call this method after adding elements to a previously empty
     * collection.
     *
     * <p>Subcollection add the ancestor's delegate instead.
     */
    void addToMap() {
      if (ancestor != null) {
        ancestor.addToMap();
      } else {
        map.put(key, delegate);
      }
    }

    @Override public int size() {
      refreshIfEmpty();
      return delegate.size();
    }

    @Override public boolean equals(Object other) {
      refreshIfEmpty();
      return delegate.equals(other);
    }

    @Override public int hashCode() {
      refreshIfEmpty();
      return delegate.hashCode();
    }

    @Override public String toString() {
      refreshIfEmpty();
      return delegate.toString();
    }

    Collection<V> getDelegate() {
      return delegate;
    }

    @Override public Iterator<V> iterator() {
      refreshIfEmpty();
      return new WrappedIterator();
    }

    /** Collection iterator for {@code WrappedCollection}. */
    class WrappedIterator implements Iterator<V> {
      final Iterator<V> delegateIterator;
      final Collection<V> originalDelegate = delegate;

      WrappedIterator() {
        delegateIterator = iteratorOrListIterator(delegate);
      }

      WrappedIterator(Iterator<V> delegateIterator) {
        this.delegateIterator = delegateIterator;
      }

      /**
       * If the delegate changed since the iterator was created, the iterator is
       * no longer valid.
       */
      void validateIterator() {
        refreshIfEmpty();
        if (delegate != originalDelegate) {
          throw new ConcurrentModificationException();
        }
      }

      public boolean hasNext() {
        validateIterator();
        return delegateIterator.hasNext();
      }

      public V next() {
        validateIterator();
        return delegateIterator.next();
      }

      public void remove() {
        delegateIterator.remove();
        totalSize--;
        removeIfEmpty();
      }

      Iterator<V> getDelegateIterator() {
        validateIterator();
        return delegateIterator;
      }
    }

    @Override public boolean add(V value) {
      refreshIfEmpty();
      boolean wasEmpty = delegate.isEmpty();
      boolean changed = delegate.add(value);
      if (changed) {
        totalSize++;
        if (wasEmpty) {
          addToMap();
        }
      }
      return changed;
    }

    WrappedCollection getAncestor() {
      return ancestor;
    }

    // The following methods are provided for better performance.

    @Override public boolean addAll(Collection<? extends V> collection) {
      if (collection.isEmpty()) {
        return false;
      }
      int oldSize = size();  // calls refreshIfEmpty
      boolean changed = delegate.addAll(collection);
      if (changed) {
        int newSize = delegate.size();
        totalSize += (newSize - oldSize);
        if (oldSize == 0) {
          addToMap();
        }
      }
      return changed;
    }

    @Override public boolean contains(Object o) {
      refreshIfEmpty();
      return delegate.contains(o);
    }

    @Override public boolean containsAll(Collection<?> c) {
      refreshIfEmpty();
      return delegate.containsAll(c);
    }

    @Override public void clear() {
      int oldSize = size();  // calls refreshIfEmpty
      if (oldSize == 0) {
        return;
      }
      delegate.clear();
      totalSize -= oldSize;
      removeIfEmpty();       // maybe shouldn't be removed if this is a sublist
    }

    @Override public boolean remove(Object o) {
      refreshIfEmpty();
      boolean changed = delegate.remove(o);
      if (changed) {
        totalSize--;
        removeIfEmpty();
      }
      return changed;
    }

    @Override public boolean removeAll(Collection<?> c) {
      if (c.isEmpty()) {
        return false;
      }
      int oldSize = size();  // calls refreshIfEmpty
      boolean changed = delegate.removeAll(c);
      if (changed) {
        int newSize = delegate.size();
        totalSize += (newSize - oldSize);
        removeIfEmpty();
      }
      return changed;
    }

    @Override public boolean retainAll(Collection<?> c) {
      int oldSize = size();  // calls refreshIfEmpty
      boolean changed = delegate.retainAll(c);
      if (changed) {
        int newSize = delegate.size();
        totalSize += (newSize - oldSize);
        removeIfEmpty();
      }
      return changed;
    }
  }

  private Iterator<V> iteratorOrListIterator(Collection<V> collection) {
    return (collection instanceof List<?>)
        ? ((List<V>) collection).listIterator()
        : collection.iterator();
  }

  /** Set decorator that stays in sync with the multimap values for a key. */
  private class WrappedSet extends WrappedCollection implements Set<V> {
    WrappedSet(@Nullable K key, Set<V> delegate) {
      super(key, delegate, null);
    }
  }

  /**
   * SortedSet decorator that stays in sync with the multimap values for a key.
   */
  private class WrappedSortedSet extends WrappedCollection
      implements SortedSet<V> {
    WrappedSortedSet(@Nullable K key, SortedSet<V> delegate,
        @Nullable WrappedCollection ancestor) {
      super(key, delegate, ancestor);
    }

    SortedSet<V> getSortedSetDelegate() {
      return (SortedSet<V>) getDelegate();
    }

    public Comparator<? super V> comparator() {
      refreshIfEmpty();
      return getSortedSetDelegate().comparator();
    }

    public V first() {
      refreshIfEmpty();
      return getSortedSetDelegate().first();
    }

    public V last() {
      refreshIfEmpty();
      return getSortedSetDelegate().last();
    }

    public SortedSet<V> headSet(V toElement) {
      refreshIfEmpty();
      return new WrappedSortedSet(
          getKey(), getSortedSetDelegate().headSet(toElement),
          (getAncestor() == null) ? this : getAncestor());
    }

    public SortedSet<V> subSet(V fromElement, V toElement) {
      refreshIfEmpty();
      return new WrappedSortedSet(
          getKey(), getSortedSetDelegate().subSet(fromElement, toElement),
          (getAncestor() == null) ? this : getAncestor());
    }

    public SortedSet<V> tailSet(V fromElement) {
      refreshIfEmpty();
      return new WrappedSortedSet(
          getKey(), getSortedSetDelegate().tailSet(fromElement),
          (getAncestor() == null) ? this : getAncestor());
    }
  }

  /** List decorator that stays in sync with the multimap values for a key. */
  private class WrappedList extends WrappedCollection
      implements List<V> {
    WrappedList(@Nullable K key, List<V> delegate,
        @Nullable WrappedCollection ancestor) {
      super(key, delegate, ancestor);
    }

    List<V> getListDelegate() {
      return (List<V>) getDelegate();
    }

    public boolean addAll(int index, Collection<? extends V> c) {
      if (c.isEmpty()) {
        return false;
      }
      int oldSize = size();  // calls refreshIfEmpty
      boolean changed = getListDelegate().addAll(index, c);
      if (changed) {
        int newSize = getDelegate().size();
        totalSize += (newSize - oldSize);
        if (oldSize == 0) {
          addToMap();
        }
      }
      return changed;
    }

    public V get(int index) {
      refreshIfEmpty();
      return getListDelegate().get(index);
    }

    public V set(int index, V element) {
      refreshIfEmpty();
      return getListDelegate().set(index, element);
    }

    public void add(int index, V element) {
      refreshIfEmpty();
      boolean wasEmpty = getListDelegate().isEmpty();
      getListDelegate().add(index, element);
      totalSize++;
      if (wasEmpty) {
        addToMap();
      }
    }

    public V remove(int index) {
      refreshIfEmpty();
      V value = getListDelegate().remove(index);
      totalSize--;
      removeIfEmpty();
      return value;
    }

    public int indexOf(Object o) {
      refreshIfEmpty();
      return getListDelegate().indexOf(o);
    }

    public int lastIndexOf(Object o) {
      refreshIfEmpty();
      return getListDelegate().lastIndexOf(o);
    }

    public ListIterator<V> listIterator() {
      refreshIfEmpty();
      return new WrappedListIterator();
    }

    public ListIterator<V> listIterator(int index) {
      refreshIfEmpty();
      return new WrappedListIterator(index);
    }

    public List<V> subList(int fromIndex, int toIndex) {
      refreshIfEmpty();
      return wrapList(getKey(), getListDelegate().subList(fromIndex, toIndex),
          (getAncestor() == null) ? this : getAncestor());
    }

    /** ListIterator decorator. */
    private class WrappedListIterator extends WrappedIterator
        implements ListIterator<V> {
      WrappedListIterator() {}

      public WrappedListIterator(int index) {
        super(getListDelegate().listIterator(index));
      }

      private ListIterator<V> getDelegateListIterator() {
        return (ListIterator<V>) getDelegateIterator();
      }

      public boolean hasPrevious() {
        return getDelegateListIterator().hasPrevious();
      }

      public V previous() {
        return getDelegateListIterator().previous();
      }

      public int nextIndex() {
        return getDelegateListIterator().nextIndex();
      }

      public int previousIndex() {
        return getDelegateListIterator().previousIndex();
      }

      public void set(V value) {
        getDelegateListIterator().set(value);
      }

      public void add(V value) {
        boolean wasEmpty = isEmpty();
        getDelegateListIterator().add(value);
        totalSize++;
        if (wasEmpty) {
          addToMap();
        }
      }
    }
  }

  /**
   * List decorator that stays in sync with the multimap values for a key and
   * supports rapid random access.
   */
  private class RandomAccessWrappedList extends WrappedList
      implements RandomAccess {
    RandomAccessWrappedList(@Nullable K key, List<V> delegate,
        @Nullable WrappedCollection ancestor) {
      super(key, delegate, ancestor);
    }
  }
  
  private transient volatile KeySet keySet;

  public Set<K> keySet() {
    if (keySet == null) {
      if (map instanceof SortedMap<?, ?>) {
        keySet = new SortedKeySet((SortedMap<K, Collection<V>>) map);
      } else {
        keySet = new KeySet(map);
      }
    }
    return keySet;
  }

  private class KeySet extends AbstractSet<K> {

    /**
     * This is usually the same as map, except when someone requests a
     * subcollection of a {@link SortedKeySet}.
     */
    final Map<K, Collection<V>> subMap;

    KeySet(final Map<K, Collection<V>> subMap) {
      this.subMap = subMap;
    }

    @Override public int size() {
      return subMap.size();
    }

    @Override public Iterator<K> iterator() {
      return new Iterator<K>() {
        final Iterator<Map.Entry<K, Collection<V>>> entryIterator
            = subMap.entrySet().iterator();
        Map.Entry<K, Collection<V>> entry;

        public boolean hasNext() {
          return entryIterator.hasNext();
        }
        public K next() {
          entry = entryIterator.next();
          return entry.getKey();
        }
        public void remove() {
          entryIterator.remove();
          totalSize -= entry.getValue().size();
          entry.getValue().clear();
        }
      };
    }

    // The following methods are included for better performance.

    @Override public boolean contains(Object key) {
      return subMap.containsKey(key);
    }

    @Override public boolean remove(Object key) {
      int count = 0;
      Collection<V> collection = subMap.remove(key);
      if (collection != null) {
        count = collection.size();
        collection.clear();
        totalSize -= count;
      }
      return count > 0;
    }

    @Override public boolean containsAll(Collection<?> c) {
      return subMap.keySet().containsAll(c);
    }
  }

  private class SortedKeySet extends KeySet implements SortedSet<K> {

    SortedKeySet(SortedMap<K, Collection<V>> subMap) {
      super(subMap);
    }

    SortedMap<K, Collection<V>> sortedMap() {
      return (SortedMap<K, Collection<V>>) subMap;
    }

    public Comparator<? super K> comparator() {
      return sortedMap().comparator();
    }

    public K first() {
      return sortedMap().firstKey();
    }

    public SortedSet<K> headSet(K toElement) {
      return new SortedKeySet(sortedMap().headMap(toElement));
    }

    public K last() {
      return sortedMap().lastKey();
    }

    public SortedSet<K> subSet(K fromElement, K toElement) {
      return new SortedKeySet(sortedMap().subMap(fromElement, toElement));
    }

    public SortedSet<K> tailSet(K fromElement) {
      return new SortedKeySet(sortedMap().tailMap(fromElement));
    }
  }
  
  private transient volatile Multiset<K> multiset;
  
  public Multiset<K> keys() {
    if (multiset == null) {
      multiset = new MultisetView();
    }
    return multiset;
  }
  
  /** Multiset view that stays in sync with the multimap keys. */
  private class MultisetView extends AbstractMultiset<K> {
    transient volatile Set<Multiset.Entry<K>> entrySet;

    @Override public int remove(Object key, int occurrences) {
      checkArgument(occurrences >= 0);
      Collection<V> collection = map.get(key);
      if (collection == null) {
        return 0;
      }
      int count = collection.size();

      if (occurrences >= count) {
        return removeValuesForKey(key);
      }

      Iterator<V> iterator = collection.iterator();
      for (int i = 0; i < occurrences; i++) {
        iterator.next();
        iterator.remove();
      }
      totalSize -= occurrences;
      return occurrences;
    }

    @Override public Set<K> elementSet() {
      return StandardMultimap.this.keySet();
    }

    @Override public Set<Multiset.Entry<K>> entrySet() {
      if (entrySet == null) {
        entrySet = new AbstractSet<Multiset.Entry<K>>() {
          @Override public Iterator<Multiset.Entry<K>> iterator() {
            return new MultisetEntryIterator();
          }
          @Override public int size() {
            return map.size();
          }

          // The following methods are included for better performance.

          @Override public boolean contains(Object o) {
            if (!(o instanceof Multiset.Entry<?>)) {
              return false;
            }
            Multiset.Entry<?> entry = (Multiset.Entry<?>) o;
            Collection<V> collection = map.get(entry.getElement());
            return (collection != null) &&
                (collection.size() == entry.getCount());
          }
          @Override public void clear() {
            StandardMultimap.this.clear();
          }
          @Override public boolean remove(Object o) {
            return contains(o) &&
                (removeValuesForKey(((Multiset.Entry<?>) o).getElement()) > 0);
          }
        };
      }
      return entrySet;
    }

    @Override public Iterator<K> iterator() {
      return new MultisetKeyIterator();
    }

    // The following methods are included for better performance.

    @Override public int count(Object key) {
      Collection<V> collection = map.get(key);
      return (collection == null) ? 0 : collection.size();
    }

    @Override public int removeAllOccurrences(Object key) {
      return removeValuesForKey(key);
    }

    @Override public int size() {
      return totalSize;
    }
  }

  /**
   * Removes all values for the provided key. Unlike {@link #removeAll}, it
   * returns the number of removed mappings.
   */
  private int removeValuesForKey(Object key) {
    int count = 0;
    Collection<V> collection = map.remove(key);
    if (collection != null) {
      count = collection.size();
      collection.clear();
      totalSize -= count;
    }
    return count;
  }

  /** Iterator across each key, repeating once per value. */
  private class MultisetEntryIterator implements Iterator<Multiset.Entry<K>> {
    final Iterator<Map.Entry<K, Collection<V>>> asMapIterator
        = asMap().entrySet().iterator();

    public boolean hasNext() {
      return asMapIterator.hasNext();
    }
    public Multiset.Entry<K> next() {
      return new MultisetEntry(asMapIterator.next());
    }
    public void remove() {
      asMapIterator.remove();
    }
  }

  private class MultisetEntry extends AbstractMultisetEntry<K> {
    final Map.Entry<K, Collection<V>> entry;

    public MultisetEntry(Map.Entry<K, Collection<V>> entry) {
      this.entry = entry;
    }

    public K getElement() {
      validate();
      return entry.getKey();
    }

    public int getCount() {
      validate();
      return entry.getValue().size();
    }

    private void validate() {
      if (entry.getValue().isEmpty()) {
        throw new ConcurrentModificationException();
      }
    }
  }

  /** Iterator across each key, repeating once per value. */
  private class MultisetKeyIterator implements Iterator<K> {
    final Iterator<Map.Entry<K, V>> entryIterator = entries().iterator();

    public boolean hasNext() {
      return entryIterator.hasNext();
    }
    public K next() {
      return entryIterator.next().getKey();
    }
    public void remove() {
      entryIterator.remove();
    }
  }
  
  private transient volatile Collection<V> values;
  
  /**
   * {@inheritDoc}
   *
   * <p>The iterator generated by the returned collection traverses the values
   * for one key, followed by the values of a second key, and so on.
   */
  public Collection<V> values() {
    if (values == null) {
      values = new AbstractCollection<V>() {
        @Override public Iterator<V> iterator() {
          return new ValueIterator();
        }
        @Override public int size() {
          return totalSize;
        }

        // The following methods are included to improve performance.

        @Override public void clear() {
          StandardMultimap.this.clear();
        }

        @Override public boolean contains(Object value) {
          return containsValue(value);
        }
      };
    }

    return values;
  }

  /** Iterator across all values. */
  private class ValueIterator implements Iterator<V> {
    final Iterator<Map.Entry<K, V>> entryIterator = createEntryIterator();

    public boolean hasNext() {
      return entryIterator.hasNext();
    }
    public V next() {
      return entryIterator.next().getValue();
    }
    public void remove() {
      entryIterator.remove();
    }
  }
  
  private transient volatile Collection<Map.Entry<K, V>> entries;
  
  /**
   * {@inheritDoc}
   *
   * <p>The iterator generated by the returned collection traverses the values
   * for one key, followed by the values of a second key, and so on.
   */
  public Collection<Map.Entry<K, V>> entries() {
    if (entries == null) {
      // TODO: can we refactor so we're not doing "this instanceof"?
      entries = (this instanceof SetMultimap<?, ?>)
          ? new EntrySet() : new Entries();
    }
    return entries;
  }

  /** Entries for multimap. */
  private class Entries extends AbstractCollection<Map.Entry<K, V>> {
    @Override public Iterator<Map.Entry<K, V>> iterator() {
      return createEntryIterator();
    }
    @Override public int size() {
      return totalSize;
    }

    // The following methods are included to improve performance.

    @Override public boolean contains(Object o) {
      if (!(o instanceof Map.Entry<?, ?>)) {
        return false;
      }
      Map.Entry<?, ?> entry = (Map.Entry<?, ?>) o;
      return containsEntry(entry.getKey(), entry.getValue());
    }

    @Override public void clear() {
      StandardMultimap.this.clear();
    }

    @Override public boolean remove(Object o) {
      if (!(o instanceof Map.Entry<?, ?>)) {
        return false;
      }
      Map.Entry<?, ?> entry = (Map.Entry<?, ?>) o;
      return StandardMultimap.this.remove(entry.getKey(), entry.getValue());
    }
  }

  /**
   * Returns an iterator across all key-value map entries, used by {@code
   * entries().iterator()} and {@code values().iterator()}. The default
   * behavior, which traverses the values for one key, the values for a second
   * key, and so on, suffices for most {@code StandardMultimap} implementations.
   *
   * @return an iterator across map entries
   */
  Iterator<Map.Entry<K, V>> createEntryIterator() {
    return new EntryIterator();
  }

  /** Iterator across all key-value pairs. */
  private class EntryIterator implements Iterator<Map.Entry<K, V>> {
    final Iterator<Map.Entry<K, Collection<V>>> keyIterator;
    K key;
    Collection<V> collection;
    Iterator<V> valueIterator;

    EntryIterator() {
      keyIterator = map.entrySet().iterator();
      if (keyIterator.hasNext()) {
        findValueIteratorAndKey();
      } else {
        valueIterator = Iterators.emptyIterator();
      }
    }

    void findValueIteratorAndKey() {
      Map.Entry<K, Collection<V>> entry = keyIterator.next();
      key = entry.getKey();
      collection = entry.getValue();
      valueIterator = iteratorOrListIterator(collection);
    }

    public boolean hasNext() {
      return valueIterator.hasNext() || keyIterator.hasNext();
    }

    public Map.Entry<K, V> next() {
      if (!hasNext()) {
        throw new NoSuchElementException();
      }
      if (!valueIterator.hasNext()) {
        findValueIteratorAndKey();
      }

      return (valueIterator instanceof ListIterator<?>)
          ? new ListMapEntry(key, (ListIterator<V>) valueIterator)
          : Maps.immutableEntry(key, valueIterator.next());
    }

    public void remove() {
      valueIterator.remove();
      if (collection.isEmpty()) {
        keyIterator.remove();
      }
      totalSize--;
    }
  }

  /**
   * Map entry that allows value modification through provided list iterator.
   */
  private class ListMapEntry extends AbstractMapEntry<K, V> {
    final K key;
    V value;
    final ListIterator<V> iterator;

    ListMapEntry(K key, ListIterator<V> iterator) {
      this.key = key;
      this.value = iterator.next();
      this.iterator = iterator;
    }
    @Override public K getKey() {
      return key;
    }
    @Override public V getValue() {
      return value;
    }
    @Override public V setValue(V newValue) {
      V oldValue = this.value;
      iterator.set(newValue);
      this.value = newValue;
      return oldValue;
    }
  }
  
  /** Entry set for a {@link SetMultimap}. */
  private class EntrySet extends Entries implements Set<Map.Entry<K, V>> {
    @Override public boolean equals(Object other) {
      if (this == other) {
        return true;
      }
      if (!(other instanceof Set<?>)) {
        return false;
      }
      Set<?> otherSet = (Set<?>) other;
      return (totalSize == otherSet.size()) && containsAll(otherSet);
    }

    @Override public int hashCode() {
      int hash = 0;
      for (Map.Entry<K, V> entry : this) {
        hash += entry.hashCode();
      }
      return hash;
    }
  }

  private transient volatile Map<K, Collection<V>> asMap;
  
  public Map<K, Collection<V>> asMap() {
    if (asMap == null) {
      asMap = new AbstractMap<K, Collection<V>>() {
        volatile Set<Map.Entry<K, Collection<V>>> entrySet;

        @Override public Set<Map.Entry<K, Collection<V>>> entrySet() {
          if (entrySet == null) {
            entrySet = new AsMapEntries();
          }
          return entrySet;
        }

        // The following methods are included for performance.

        @Override public void clear() {
          StandardMultimap.this.clear();
        }

        @Override public boolean containsKey(Object key) {
          return map.containsKey(key);
        }

        @Override public Collection<V> get(Object key) {
          Collection<V> collection = map.get(key);
          if (collection == null) {
            return null;
          }
          @SuppressWarnings("unchecked")
          K k = (K) key;
          return wrapCollection(k, collection);
        }

        @Override public Collection<V> remove(Object key) {
          Collection<V> collection = removeAll(key);
          return collection.isEmpty() ? null : collection;
        }

        @Override public boolean equals(Object other) {
          return map.equals(other);
        }

        @Override public int hashCode() {
          return map.hashCode();
        }

        @Override public String toString() {
          return map.toString();
        }
      };
    }

    return asMap;
  }

  private class AsMapEntries extends AbstractSet<Map.Entry<K, Collection<V>>> {
    @Override public Iterator<Map.Entry<K, Collection<V>>> iterator() {
      return new AsMapIterator();
    }

    @Override public int size() {
      return map.size();
    }

    // The following methods are included for performance.

    @Override public void clear() {
      StandardMultimap.this.clear();
    }

    @Override public boolean contains(Object o) {
      if (!(o instanceof Map.Entry<?, ?>)) {
        return false;
      }
      Map.Entry<?, ?> entry = (Map.Entry<?, ?>) o;
      return (entry.getValue() != null)
          && Objects.equal(map.get(entry.getKey()), entry.getValue());
    }

    @Override public boolean remove(Object o) {
      if (!(o instanceof Map.Entry<?, ?>)) {
        return false;
      }
      Map.Entry<?, ?> entry = (Map.Entry<?, ?>) o;
      return contains(entry) && (removeValuesForKey(entry.getKey()) > 0);
    }
  }

  /** Iterator across all keys and value collections. */
  private class AsMapIterator implements Iterator<Map.Entry<K, Collection<V>>> {
    final Iterator<Map.Entry<K, Collection<V>>> delegateIterator
        = map.entrySet().iterator();
    Collection<V> collection;

    public boolean hasNext() {
      return delegateIterator.hasNext();
    }
    public Map.Entry<K, Collection<V>> next() {
      Map.Entry<K, Collection<V>> entry = delegateIterator.next();
      K key = entry.getKey();
      collection = entry.getValue();
      return Maps.immutableEntry(key, wrapCollection(key, collection));
    }
    public void remove() {
      delegateIterator.remove();
      totalSize -= collection.size();
      collection.clear();
    }
  }

  // Comparison and hashing
  
  @Override public boolean equals(@Nullable Object other) {
    if (this == other) {
      return true;
    }
    if (!(other instanceof Multimap<?, ?>)) {
      return false;
    }
    Multimap<?, ?> otherMultimap = (Multimap<?, ?>) other;
    return map.equals(otherMultimap.asMap());
  }

  /**
   * Returns the hash code for this multimap.
   *
   * <p>The hash code of a multimap is defined as the hash code of the map view,
   * as returned by {@link Multimap#asMap}.
   *
   * @see Map#hashCode
   */
  @Override public int hashCode() {
    return map.hashCode();
  }

  /**
   * Returns a string representation of the multimap, generated by calling
   * {@code toString} on the map returned by {@link Multimap#asMap}.
   *
   * @return a string representation of the multimap
   */
  @Override public String toString() {
    return map.toString();
  }
}