
package de.ivu.elect.util;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import java.util.function.Function;

public class CollectionSeparator<Element, ElementKey> {

  private final Collection<Element> first;

  private final Collection<Element> second;

  private final Function<Element, ElementKey> keyFunction;

  private  Collection<Element> inFirstNotInSecond = null;

  private  Collection<Element> inSecondNotInFirst = null;

  private  Collection<ElementKey> inBoth = null;

  private  Map<ElementKey, Element> firstMap = null;

  private  Map<ElementKey, Element> secondMap = null;

  public CollectionSeparator(Collection<Element> first, Collection<Element> second,
                             Function<Element, ElementKey> keyFunction) {
    this.first = first;
    this.second = second;
    this.keyFunction = keyFunction;
  }

  private void lazy(){
    if (inFirstNotInSecond==null){
      inFirstNotInSecond= new ArrayList<>(first.size());
      inSecondNotInFirst = new ArrayList<>(second.size());
      inBoth = new ArrayList<>(Math.min(first.size(), second.size()));
      firstMap = MapBuilder.newMapFrom(first, keyFunction);
      secondMap = MapBuilder.newMapFrom(second, keyFunction);

      for (Map.Entry<ElementKey, Element> firstEntry : firstMap.entrySet()) {
        if (secondMap.containsKey(firstEntry.getKey())) {
          inBoth.add(firstEntry.getKey());
        } else {
          inFirstNotInSecond.add(firstEntry.getValue());
        }
      }

      for (Map.Entry<ElementKey, Element> secondEntry : secondMap.entrySet()) {
        if (!firstMap.containsKey(secondEntry.getKey())) {
          inSecondNotInFirst.add(secondEntry.getValue());
        }
      }
    }
  }

  public Collection<Element> getInFirstNotInSecond() {
    if (first.isEmpty() || second.isEmpty()){
      return first;
    }

    lazy();
    return inFirstNotInSecond;
  }

  public Collection<Element> getInSecondNotInFirst() {
    if (first.isEmpty() || second.isEmpty()){
      return second;
    }

    lazy();
    return inSecondNotInFirst;
  }

  public Collection<ElementKey> getInBoth() {
    if (first.isEmpty() || second.isEmpty()){
      return Collections.emptyList();
    }
    lazy();
    return inBoth;
  }

  public Map<ElementKey, Element> getFirstMap() {
    if (first.isEmpty()){
      return new HashMap<>();
    }
    lazy();
    return firstMap;
  }

  public Map<ElementKey, Element> getSecondMap() {
    if(second.isEmpty()){
      return new HashMap<>();
    }
    lazy();
    return secondMap;
  }
}
