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

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

import org.apache.uima.jcas.JCas;

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

import uk.gov.dstl.baleen.annotators.coreference.impl.data.Animacy;
import uk.gov.dstl.baleen.annotators.coreference.impl.data.Cluster;
import uk.gov.dstl.baleen.annotators.coreference.impl.data.Mention;
import uk.gov.dstl.baleen.annotators.coreference.impl.data.MentionType;
import uk.gov.dstl.baleen.annotators.coreference.impl.data.Person;
import uk.gov.dstl.baleen.resources.data.Gender;

/**
 * Attempts to connect pronouns to an entity.
 *
 * <p>This is very difficult problem which may perform suboptimally in the current implementation.
 */
public class PronounResolutionSieve extends AbstractCoreferenceSieve {

  private static final int MAX_SENTENCE_DISTANCE = 3;

  /** Constructor for PronounResolutionSieve */
  public PronounResolutionSieve(JCas jCas, List<Cluster> clusters, List<Mention> mentions) {
    super(jCas, clusters, mentions);
  }

  @Override
  public void sieve() {

    // Find potential entity->pronoun matches
    final Multimap<Mention, Mention> potential = HashMultimap.create();

    // Check that all mentions have a valid sentence index
    List<Mention> filteredMentions =
        getMentions().stream().filter(m -> m.getSentenceIndex() >= 0).collect(Collectors.toList());

    for (int i = 0; i < filteredMentions.size(); i++) {
      final Mention a = filteredMentions.get(i);

      for (int j = i + 1; j < filteredMentions.size(); j++) {
        final Mention b = filteredMentions.get(j);

        if (!validateMentions(a, b)) {
          continue;
        }

        resolvePronoun(a, b, potential);
      }
    }

    // For each of the matches we need to select the best one

    potential.asMap().entrySet().stream().forEach(e -> addBestAsMatch(e.getKey(), e.getValue()));
  }

  /** If the pronoun meets the conditions in this function, it is *likely* to be a person */
  private boolean isPronounPerson(Mention pronoun) {
    return pronoun.getPerson() != Person.UNKNOWN // Acceptable values: FIRST, SECOND and THIRD
            && (pronoun.getGender() == Gender.M || pronoun.getGender() == Gender.F)
        || pronoun.getAnimacy() == Animacy.ANIMATE;
  }

  private boolean validateMentions(Mention a, Mention b) {
    // We are coreferencing pronouns only, so a OR b must be a pronoun, but not both
    if (!(a.getType() == MentionType.PRONOUN && b.getType() != MentionType.PRONOUN
        || a.getType() != MentionType.PRONOUN && b.getType() == MentionType.PRONOUN)) {
      return false;
    }

    // Not in paper: No overlap it makes little sense
    if (a.overlaps(b)) {
      return false;
    }

    // Are the attributes compatible (gender=gender, etc)
    return a.isAttributeCompatible(b);
  }

  private void resolvePronoun(Mention a, Mention b, Multimap<Mention, Mention> potential) {
    final Mention pronoun = a.getType() == MentionType.PRONOUN ? a : b;
    final Mention other = a.getType() == MentionType.PRONOUN ? b : a;

    // Not in paper: If the pronoun is before the other that's odd, (He said Hello. John
    // did.)
    if (pronoun.getAnnotation().getEnd() < other.getAnnotation().getBegin()) {
      return;
    }

    // Not in paper: We found poor results for "it" really because it never refers to something that
    // Baleen annotates
    // It would be good for say money but currently we'll just drop it
    if (pronoun.getText().toLowerCase().startsWith("it")) {
      return;
    }

    // Originally there was code here to handle a and b both being pronouns (this is described in
    // the paper)
    // However, we are excluding links between two pronouns in validateMentions, so this code was
    // removed.

    // Paper: Only consider within three
    // Not in paper: And the pronoun must be after the mention
    final int sentenceDistance = pronoun.getSentenceIndex() - other.getSentenceIndex();
    if (sentenceDistance < 0 || sentenceDistance > MAX_SENTENCE_DISTANCE) {
      return;
    }

    // Not in paper: If the same sentence the pronoun should be after
    if (sentenceDistance == 0
        && pronoun.getAnnotation().getEnd() <= other.getAnnotation().getBegin()) {
      return;
    }

    if (isPronounPerson(pronoun)
        && !(other.getAnnotation() instanceof uk.gov.dstl.baleen.types.common.Person
            || other.getAnnotation() instanceof uk.gov.dstl.baleen.types.common.Nationality)) {
      return;
    }

    // Similarly to avoid if we have a neutral we can't link to a person
    if (pronoun.getGender() == Gender.N
        && other.getAnnotation() instanceof uk.gov.dstl.baleen.types.common.Person) {
      return;
    }

    // TODO: There might be many more of these simple constraints on our semantic types...

    potential.put(pronoun, other);
  }

  private void addBestAsMatch(Mention key, Collection<Mention> potentialMatches) {

    final Collection<Mention> matched;
    if (potentialMatches.size() > 1) {
      List<Mention> list = new ArrayList<>(potentialMatches);
      Collections.sort(
          list,
          (a, b) -> {
            if (a.overlaps(b)) {
              return 0;
            }

            // Use in-sentence word distance
            if (a.getAnnotation().getEnd() <= b.getAnnotation().getBegin()) {
              return b.getAnnotation().getBegin() - a.getAnnotation().getEnd();
            } else {
              return b.getAnnotation().getEnd() - a.getAnnotation().getBegin();
            }
          });

      matched = list;
    } else {
      // Either empty or just one...
      matched = potentialMatches;
    }

    // Get the first (nearest) which doesn't overlap
    Optional<Mention> match = matched.stream().filter(m -> !key.overlaps(m)).findFirst();

    if (match.isPresent()) {
      addToCluster(key, match.get());
    }
  }
}
