// Dstl (c) Crown Copyright 2017
package uk.gov.dstl.baleen.annotators.cleaners.helpers;

import java.util.Collection;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import org.apache.uima.analysis_engine.AnalysisEngineProcessException;
import org.apache.uima.jcas.JCas;

import uk.gov.dstl.baleen.types.semantic.Entity;
import uk.gov.dstl.baleen.uima.BaleenAnnotator;

/**
 * <b>Remove entities which are contained within other entities of the same type</b>
 *
 * <p>All entities are looped through, and should an entity be found to be entirely contained within
 * another entity of the same type it is removed. The comparison is done purely on start and end
 * positions, and ignores other information within the entity. If two entities of the same type have
 * the same start and end position, then the one with the lower confidence is removed; and if both
 * have the same confidence then the first entity in the annotation index is removed.
 *
 * <p>Information may be lost by the removal of entities. Some kind of merging of entities might be
 * a better option.
 *
 * <p>Implementations should override shouldMerge to determine if item should be merged, and the
 * compileEntities to create lists of entities to be considered for merging.
 */
public abstract class AbstractNestedEntities<T extends Entity> extends BaleenAnnotator {

  @Override
  public void doProcess(JCas jCas) throws AnalysisEngineProcessException {

    Collection<List<T>> annotations = compileEntities(jCas);

    List<MergePair<T>> mergePairs = new LinkedList<>();
    Set<T> toRemove = new HashSet<>();

    for (List<T> entities : annotations) {
      removeNestedEntities(entities, mergePairs, toRemove);
    }

    for (MergePair<T> pair : mergePairs) {
      T keep = pair.getKeep();
      T remove = pair.getRemove();
      if (shouldMerge(keep, remove)) {
        // We used to log the change here, but calling
        // having any of the parameters in this debug causes a javac compile error
        // Seems to be a recognised JDK see JDK-8056984, to be fixed in JDK 9
        // Issue arose when introducing History. Perhaps this has to do
        // with the Recordable interface?

        mergeWithExisting(keep, remove);
      } else {
        // If we aren't merging this we should keep both
        toRemove.remove(remove);
      }
    }
    removeFromJCasIndex(toRemove);
  }

  private void removeNestedEntities(
      List<T> typeAnnotations, List<MergePair<T>> mergePairs, Set<T> toRemove) {
    for (int x = 0; x < typeAnnotations.size(); x++) {
      T eX = typeAnnotations.get(x);

      for (int y = x + 1; y < typeAnnotations.size(); y++) {
        T eY = typeAnnotations.get(y);

        if (toRemove.contains(eX) || toRemove.contains(eY)) {
          // If we are always removing the one of the entities, we should have already
          // dealt with other cases
          continue;
        }

        if (eX.getBegin() == eY.getBegin() && eX.getEnd() == eY.getEnd()) {
          removeLeastConfidentEntity(mergePairs, toRemove, eX, eY);
        } else if (containedWithin(eX, eY)) {
          // Remove nested entity Y
          removeEntity(mergePairs, toRemove, eY, eX);
        } else if (containedWithin(eY, eX)) {
          // Remove nested entity Y
          removeEntity(mergePairs, toRemove, eX, eY);
        }
      }
    }
  }

  /** Is e2 contained within e1? */
  private boolean containedWithin(T e1, T e2) {
    return e1.getBegin() <= e2.getBegin() && e1.getEnd() >= e2.getEnd();
  }

  /** Removes the least confident entity, or e2 if they have the same confidence */
  private void removeLeastConfidentEntity(
      List<MergePair<T>> mergePairs, Set<T> toRemove, T e1, T e2) {
    if (e1.getConfidence() >= e2.getConfidence()) {
      removeEntity(mergePairs, toRemove, e2, e1);
    } else {
      removeEntity(mergePairs, toRemove, e1, e2);
    }
  }

  private void removeEntity(
      List<MergePair<T>> mergePairs, Set<T> toRemove, T entityToRemove, T entityToKeep) {
    toRemove.add(entityToRemove);
    mergePairs.add(new MergePair<T>(entityToKeep, entityToRemove));
  }

  /**
   * Get list of entities to processed for overlap and nesting (and hence merged together), each
   * list will be treated independently.
   *
   * @param jCas the CAS to pull entities from.
   * @return a collection of lists. The lists should contain entities to be compared together.
   *     typically this will be entities of the same type.
   */
  protected abstract Collection<List<T>> compileEntities(JCas jCas);

  /**
   * Determine if these specific pairs of entities should be merged together.
   *
   * @param keep the entity to be kept
   * @param remove the entity being considered for removal.
   * @return false if both entities should be kept, true is the remove entity should be merged into
   *     the kept.
   */
  protected abstract boolean shouldMerge(T keep, T remove);
}
