package com.bbn.bue.common.collections;

import com.google.common.annotations.Beta;
import com.google.common.base.Function;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSetMultimap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Multimap;

import java.util.Collection;
import java.util.Map;

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

public final class MultimapUtils {

  private MultimapUtils() {
    throw new UnsupportedOperationException();
  }

  /**
   * This will take a multimap which in fact has a single value map for each key and return a copy
   * of it as an {@link com.google.common.collect.ImmutableMap}. If the same key if mapped to
   * multiple values in the input multimap, an {@link java.lang.IllegalArgumentException} is thrown,
   * so this should generally not be used on multimaps built from user input unless this property is
   * guaranteed to hold.  This is mostly useful for statically building maps which are more
   * naturally represented as inverted multimaps.
   *
   * Preserves iteration order.
   */
  public static <K, V> Map<K, V> copyAsMap(Multimap<K, V> multimap) {
    final Map<K, Collection<V>> inputAsMap = multimap.asMap();
    final ImmutableMap.Builder<K, V> ret = ImmutableMap.builder();

    for (final Map.Entry<K, Collection<V>> mapping : inputAsMap.entrySet()) {
      ret.put(mapping.getKey(), Iterables.getOnlyElement(mapping.getValue()));
    }

    return ret.build();
  }

  /**
   * Converts a {@link com.google.common.collect.Multimap} to a {@link java.util.Map} by selecting a
   * single value for each key and discarding the others, with the selection procedure being defined
   * by {@code reducerFunction}, which may not return {@code null}. Additionally, {@code multimap}
   * may not have {@code null} keys.
   */
  public static <K, V> ImmutableMap<K, V> reduceToMap(Multimap<K, V> multimap,
      Function<? super Collection<V>, ? extends V> reducerFunction) {
    final ImmutableMap.Builder<K, V> ret = ImmutableMap.builder();

    for (final Map.Entry<K, Collection<V>> entry : multimap.asMap().entrySet()) {
      ret.put(entry.getKey(), reducerFunction.apply(entry.getValue()));
    }

    return ret.build();
  }

  /**
   * Performs a (non-strict) composition of two multimaps to an {@link ImmutableSetMultimap}.
   * This returns a new {@link ImmutableSetMultimap} which will contain an entry {@code (k,v)} if
   * and only if there is some {@code i} such that {@code (k, i)} is in {@code first}
   * and {@code (i,k)} is in {@code second}.  Neither {@code first} nor {@code second} is permitted
   * to contain null keys or values. The output of this method is not a view and will not be updated
   * for changes to the input multimaps.
   *
   * This method will allow {@code first} to contain values not found as keys of {@code second}. If
   * you wish to disallow this, see {@link #composeToSetMultimapStrictly(Multimap, Multimap)}.
   */
  @Beta
  public static <K1, K2, V1 extends K2, V2> ImmutableSetMultimap<K1, V2> composeToSetMultimap(
      final Multimap<K1, V1> first, final Multimap<K2, V2> second) {
    final ImmutableSetMultimap.Builder<K1, V2> result = ImmutableSetMultimap.builder();
    for (K1 k1 : first.keySet()) {
      checkNotNull(k1);
      for (V1 v1 : first.get(k1)) {
        checkNotNull(v1);
        result.putAll(k1, second.get(v1));
      }
    }
    return result.build();
  }

  /**
   * Performs a composition on two multimaps, strictly asserting that the domain of the second
   * contains the range of the first.
   *
   * Performs the same operation as {@link MultimapUtils#composeToSetMultimap(Multimap, Multimap)},
   * ensuring that the sets V1 and K2 are identical.
   */
  @Beta
  public static <K1, K2, V1 extends K2, V2> ImmutableSetMultimap<K1, V2> composeToSetMultimapStrictly(
      final Multimap<K1, V1> first, final Multimap<K2, V2> second) {
    checkArgument(second.keySet().containsAll(first.values()));
    return composeToSetMultimap(first, second);
  }
}

@@ -2443,26 +2443,26 @@
 ltimap%7D.
-%0A   *
  This
+%0A   *
  returns
@@ -2541,25 +2541,25 @@
 k,v)%7D if
+ and
 %0A   *
- and
  only if
@@ -2627,21 +2627,16 @@
 e first%7D
-%0A   *
  and %7B@c
@@ -2646,11 +2646,16 @@
  (i,
-k
+v
 )%7D
+%0A   *
  is 
@@ -2727,21 +2727,16 @@
 ermitted
-%0A   *
  to cont
@@ -2743,16 +2743,21 @@
 ain null
+%0A   *
  keys or
@@ -2828,21 +2828,16 @@
  updated
-%0A   *
  for cha
@@ -2843,16 +2843,21 @@
 anges to
+%0A   *
  the inp
@@ -3567,16 +3567,25 @@
 rforms a
+ (strict)
  composi
@@ -3590,17 +3590,17 @@
 sition o
-n
+f
  two mul
@@ -3609,253 +3609,610 @@
 maps
-, strictly asserting that the domain of the second%0A   * contains the range of the first.%0A   *%0A   * Performs the same operation as %7B@link MultimapUtils#composeToSetMultimap(Multimap, Multimap)%7D,%0A   * ensuring that the sets V1 and K2 are identical
+ to an %7B@link ImmutableSetMultimap%7D. This%0A   * returns a new %7B@link ImmutableSetMultimap%7D which will contain an entry %7B@code (k,v)%7D if and%0A   * only if there is some %7B@code i%7D such that %7B@code (k, i)%7D is in %7B@code first%7D and %7B@code (i, v)%7D%0A   * is in %7B@code second%7D.  Neither %7B@code first%7D nor %7B@code second%7D is permitted to contain null%0A   * keys or values. The output of this method is not a view and will not be updated for changes to%0A   * the input multimaps. Strict compositions require that for each entry %7B@code (k,i)%7D in %7B@code%0A   * first%7D, there exists an entry %7B@code (i,v)%7D in %7B@code second%7D
 .%0A  

change composeToSetMultimapStrictly to be consistent with non-strict version