package com.bbn.bue.common.evaluation;

import com.google.common.annotations.Beta;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.ImmutableSetMultimap;
import com.google.common.collect.Multimap;

import java.util.Collection;
import java.util.Set;

import static com.google.common.collect.Sets.difference;

/**
 * A {@link ProvenancedAlignment} based on grouping items into equivalence classes. The items
 * aligned are the equivalence classes and the provenances are the original items.
 */
@Beta
public final class EquivalenceBasedProvenancedAlignment<EqClassT, LeftT, RightT>
    implements ProvenancedAlignment<EqClassT, LeftT, EqClassT, RightT> {

  private final ImmutableMultimap<EqClassT, LeftT> leftEquivalenceClassesToProvenances;
  private final ImmutableMultimap<EqClassT, RightT> rightEquivalenceClassesToProvenances;

  private EquivalenceBasedProvenancedAlignment(
      final Multimap<? extends EqClassT, ? extends LeftT> leftEquivalenceClassesToProvenances,
      final Multimap<? extends EqClassT, ? extends RightT> rightEquivalenceClassesToProvenances) {
    this.leftEquivalenceClassesToProvenances = ImmutableSetMultimap
        .copyOf(leftEquivalenceClassesToProvenances);
    this.rightEquivalenceClassesToProvenances = ImmutableSetMultimap.copyOf(
        rightEquivalenceClassesToProvenances);
  }

  // package-private
  static <EqClassT, LeftProvT, RightProvT> EquivalenceBasedProvenancedAlignment<EqClassT, LeftProvT, RightProvT> fromEquivalenceClassMaps(
      final Multimap<? extends EqClassT, ? extends LeftProvT> leftEquivalenceClassesToProvenances,
      final Multimap<? extends EqClassT, ? extends RightProvT> rightEquivalenceClassesToProvenances) {
    return new EquivalenceBasedProvenancedAlignment<EqClassT, LeftProvT, RightProvT>(
        leftEquivalenceClassesToProvenances, rightEquivalenceClassesToProvenances);
  }

  @Override
  public Collection<LeftT> provenancesForLeftItem(final EqClassT item) {
    return leftEquivalenceClassesToProvenances.get(item);
  }

  @Override
  public Collection<RightT> provenancesForRightItem(final EqClassT item) {
    return rightEquivalenceClassesToProvenances.get(item);
  }

  @Override
  public Set<EqClassT> leftUnaligned() {
    return difference(leftEquivalenceClassesToProvenances.keySet(),
        rightEquivalenceClassesToProvenances.keySet());
  }

  @Override
  public Set<EqClassT> rightUnaligned() {
    return difference(rightEquivalenceClassesToProvenances.keySet(),
        leftEquivalenceClassesToProvenances.keySet());
  }

  @Override
  public Set<EqClassT> leftAligned() {
    return leftEquivalenceClassesToProvenances.keySet();
  }

  @Override
  public Set<EqClassT> rightAligned() {
    return rightEquivalenceClassesToProvenances.keySet();
  }

  // if it appears in the multimap, it's got to be an EqClassT
  @SuppressWarnings("unchecked")
  @Override
  public Collection<EqClassT> alignedToRightItem(final Object rightItem) {
    if (rightEquivalenceClassesToProvenances.containsKey(rightItem)) {
      return ImmutableList.of((EqClassT) rightItem);
    } else {
      return ImmutableList.of();
    }
  }

  // if it appears in the multimap, it's got to be an EqClassT
  @SuppressWarnings("unchecked")
  @Override
  public Collection<EqClassT> alignedToLeftItem(final Object leftItem) {
    if (rightEquivalenceClassesToProvenances.containsKey(leftItem)) {
      return ImmutableList.of((EqClassT) leftItem);
    } else {
      return ImmutableList.of();
    }
  }

  @Override
  public Set<EqClassT> allLeftItems() {
    return leftEquivalenceClassesToProvenances.keySet();
  }

  @Override
  public Set<EqClassT> allRightItems() {
    return rightEquivalenceClassesToProvenances.keySet();
  }
}

@@ -3023,179 +3023,39 @@
     
-if (rightEquivalenceClassesToProvenances.containsKey(rightItem)) %7B%0A      return ImmutableList.of((EqClassT) rightItem);%0A    %7D else %7B%0A      return ImmutableList.of();%0A    %7D
+return getAlignedTo(rightItem);
 %0A  %7D
@@ -3245,113 +3245,77 @@
     
-if (rightEquivalenceClassesToProvenances.containsKey(leftItem)) %7B%0A      return ImmutableList.of((
+return getAlignedTo(leftItem);%0A  %7D%0A%0A  @Override%0A  public Set%3C
 EqClassT
 ) le
@@ -3314,36 +3314,27 @@
 assT
-) l
+%3E allL
 eftItem
-);%0A    %7D else
+s()
  %7B%0A
-  
     
@@ -3344,33 +3344,53 @@
 urn 
-ImmutableList.of();%0A    %7D
+leftEquivalenceClassesToProvenances.keySet();
 %0A  %7D
@@ -3421,35 +3421,36 @@
 et%3CEqClassT%3E all
-Lef
+Righ
 tItems() %7B%0A    r
@@ -3447,35 +3447,36 @@
 () %7B%0A    return 
-lef
+righ
 tEquivalenceClas
@@ -3513,53 +3513,232 @@
 %0A%0A  
-@Override%0A  public Set%3CEqClassT%3E allR
+/**%0A   * Any equivalence class is by definition aligned to itself if it is present on both the left%0A   * and the right. Otherwise, it has no alignment.%0A   */%0A  private Collection%3CEqClassT%3E getAlignedTo(final Object r
 ightItem
 s() 
@@ -3729,34 +3729,32 @@
 ct rightItem
-s(
 ) %7B%0A    
 return right
@@ -3745,61 +3745,249 @@
     
-return rightEquivalenceClassesToProvenances.keySet();
+if (rightEquivalenceClassesToProvenances.containsKey(rightItem)%0A        && leftEquivalenceClassesToProvenances.containsKey(rightItem)) %7B%0A      return ImmutableList.of((EqClassT) rightItem);%0A    %7D else %7B%0A      return ImmutableList.of();%0A    %7D
 %0A  %7D

Fix #7 by verifying item is in left and right.