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

import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

import org.apache.uima.analysis_engine.AnalysisEngineProcessException;
import org.apache.uima.fit.util.JCasUtil;
import org.apache.uima.jcas.JCas;

import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Multimap;

import uk.gov.dstl.baleen.core.pipelines.orderers.AnalysisEngineAction;
import uk.gov.dstl.baleen.types.semantic.Entity;
import uk.gov.dstl.baleen.types.semantic.ReferenceTarget;
import uk.gov.dstl.baleen.uima.BaleenAnnotator;
import uk.gov.dstl.baleen.uima.utils.ComparableEntitySpanUtils;
import uk.gov.dstl.baleen.uima.utils.ReferentUtils;

/**
 * Convert non-entity annotations into entity annotations, in the case that an annotation has a
 * referent target that is shared with an entity.
 *
 * <p>This is useful for consumers that work specifically with entities but not with other types.
 *
 * @baleen.javadoc
 */
public class ReferentToEntity extends BaleenAnnotator {

  @Override
  @SuppressWarnings("squid:S2175" /* false positive, a Base can be an Entity */)
  protected void doProcess(JCas jCas) throws AnalysisEngineProcessException {

    final Multimap<ReferenceTarget, Entity> referentMap =
        ReferentUtils.createReferentMap(jCas, Entity.class);

    final Collection<Entity> entities = new HashSet<>(JCasUtil.select(jCas, Entity.class));

    final Map<ReferenceTarget, Entity> targets =
        ReferentUtils.filterToSingle(referentMap, ReferentToEntity::getBestEntity);

    // Now look through the non-entities and create entities in their place.

    final List<Entity> toAdd =
        ReferentUtils.streamReferent(jCas, targets)
            .map(
                a -> {
                  final ReferenceTarget referent = a.getReferent();
                  final Entity entity = targets.get(referent);
                  if (entity != null && !entities.contains(a)) {
                    return ComparableEntitySpanUtils.copyEntity(
                        jCas, a.getBegin(), a.getEnd(), entity);
                  } else {
                    return null;
                  }
                })
            .filter(Objects::nonNull)
            .collect(Collectors.toList());

    addToJCasIndex(toAdd);
  }

  /**
   * Gets the best entity from the list.
   *
   * @param list the list
   * @return the best entity
   */
  protected static Entity getBestEntity(Collection<Entity> list) {
    return list.stream()
        .reduce((a, b) -> isBetterEntity(a, b) ? b : a)
        .orElseThrow(() -> new RuntimeException("Should always produce an entity"));
  }

  /**
   * Checks if is better entity.
   *
   * @param original the original
   * @param challenger the challenger
   * @return true, if is better entity
   */
  protected static boolean isBetterEntity(Entity original, Entity challenger) {
    // Simple version, just look for the longest string
    // we could look at how complete the attributes are, etc
    String origValue = original.getValue();
    if (origValue == null) {
      origValue = original.getCoveredText();
    }

    String challValue = challenger.getValue();
    if (challValue == null) {
      challValue = challenger.getCoveredText();
    }

    return origValue.length() < challValue.length();
  }

  @Override
  public AnalysisEngineAction getAction() {
    return new AnalysisEngineAction(ImmutableSet.of(Entity.class), ImmutableSet.of(Entity.class));
  }
}
