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

import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;

import org.apache.uima.jcas.JCas;

import com.google.common.base.Strings;

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.types.language.PhraseChunk;

/**
 * Base class for coreference sieves.
 *
 * <p>Provides helper functions to manage clusters, etc.
 */
public abstract class AbstractCoreferenceSieve implements CoreferenceSieve {

  private final JCas jCas;
  private final List<Cluster> clusters;
  private final List<Mention> mentions;

  protected AbstractCoreferenceSieve(JCas jCas, List<Cluster> clusters, List<Mention> mentions) {
    this.jCas = jCas;
    this.clusters = clusters;
    this.mentions = mentions;
  }

  protected List<Cluster> getClusters() {
    return clusters;
  }

  protected List<Mention> getMentions() {
    return mentions;
  }

  protected List<Mention> getMentions(MentionType mentionType) {
    return mentions.stream()
        .filter(m -> m.getType().equals(mentionType))
        .collect(Collectors.toList());
  }

  protected List<Mention> getMentions(MentionType... mentionTypes) {
    List<MentionType> mentionTypesList = Arrays.asList(mentionTypes);
    return mentions.stream()
        .filter(m -> mentionTypesList.contains(m.getType()))
        .collect(Collectors.toList());
  }

  protected List<Mention> getMentionsWithHead(MentionType mentionType) {
    return mentions.stream()
        .filter(m -> m.getType().equals(mentionType) && hasHead(m))
        .collect(Collectors.toList());
  }

  protected List<Mention> getMentionsWithHead(MentionType... mentionTypes) {
    List<MentionType> mentionTypesList = Arrays.asList(mentionTypes);

    return mentions.stream()
        .filter(m -> mentionTypesList.contains(m.getType()) && hasHead(m))
        .collect(Collectors.toList());
  }

  protected JCas getJCas() {
    return jCas;
  }

  protected void addToCluster(Mention a, Mention b) {
    if (a.hasClusters()) {
      final Cluster cluster = a.getAnyCluster();
      cluster.add(b);
    } else if (b.hasClusters()) {
      final Cluster cluster = b.getAnyCluster();
      cluster.add(a);
    } else {
      clusters.add(new Cluster(a, b));
    }
  }

  protected void addPairwiseToCluster(Collection<Mention> a, Collection<Mention> b) {

    // Technically these will all need to end up in the same cluster

    final Cluster cluster =
        a.stream()
            .map(Mention::getAnyCluster)
            .filter(Objects::nonNull)
            .findAny()
            .orElseGet(
                () ->
                    b.stream()
                        .map(Mention::getAnyCluster)
                        .filter(Objects::nonNull)
                        .findAny()
                        .orElseGet(
                            () -> {
                              final Cluster c = new Cluster();
                              clusters.add(c);
                              return c;
                            }));

    cluster.addAll(a);
    cluster.addAll(b);
  }

  protected void addCoveredToCluster(PhraseChunk a, PhraseChunk b) {
    final List<Mention> aMentions = findMentionsExactly(a.getBegin(), a.getEnd());
    final List<Mention> bMentions = findMentionsExactly(b.getBegin(), b.getEnd());

    addPairwiseToCluster(aMentions, bMentions);
  }

  protected List<Mention> findMentionsExactly(int begin, int end) {
    return getMentions().stream()
        .filter(m -> begin == m.getAnnotation().getBegin() && m.getAnnotation().getEnd() == end)
        .collect(Collectors.toList());
  }

  protected List<Mention> findMentionsUnder(int begin, int end) {
    return getMentions().stream()
        .filter(m -> begin >= m.getAnnotation().getBegin() && m.getAnnotation().getEnd() <= end)
        .collect(Collectors.toList());
  }

  protected List<Mention> findMentionAbove(int begin, int end) {
    return getMentions().stream()
        .filter(m -> m.getAnnotation().getBegin() <= begin && end <= m.getAnnotation().getEnd())
        .collect(Collectors.toList());
  }

  protected Set<String> getModifiers(Mention a) {
    // In the paper they say N and J (adjective) but we need cardinal too otherwise "200 people"
    // discards the 200
    // TODO: Modifiers up to head word? See paper
    return a.getWords().stream()
        .filter(
            w ->
                w.getPartOfSpeech().startsWith("N")
                    || w.getPartOfSpeech().startsWith("J")
                    || w.getPartOfSpeech().startsWith("CD"))
        .map(w -> w.getCoveredText().toLowerCase())
        .collect(Collectors.toSet());
  }

  protected static boolean hasHead(Mention m) {
    return !Strings.isNullOrEmpty(m.getHead());
  }
}
