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

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

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

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;
import uk.gov.dstl.baleen.types.language.Sentence;
import uk.gov.dstl.baleen.types.language.WordToken;
import uk.gov.dstl.baleen.types.semantic.Entity;

/**
 * Adds the sentence and its index (count from the start) to each mention.
 *
 * <p>Thanks to UIMA this is a tremendous amount of work for such a simple task.
 */
public class SentenceEnhancer {

  /**
   * Enhance the mentions by adding sentence information.
   *
   * @param jCas the j cas
   * @param mentions the mentions
   */
  public void enhance(JCas jCas, List<Mention> mentions) {
    // Create a map (mention annotation) to sentence

    final Set<WordToken> pronounAnnotation =
        mentions.stream()
            .filter(p -> p.getType() == MentionType.PRONOUN)
            .map(p -> (WordToken) p.getAnnotation())
            .collect(Collectors.toSet());

    final Set<Entity> entityAnnotation =
        mentions.stream()
            .filter(p -> p.getType() == MentionType.ENTITY)
            .map(p -> (Entity) p.getAnnotation())
            .collect(Collectors.toSet());

    final Set<PhraseChunk> npAnnotation =
        mentions.stream()
            .filter(p -> p.getType() == MentionType.NP)
            .map(p -> (PhraseChunk) p.getAnnotation())
            .collect(Collectors.toSet());

    final Map<WordToken, Collection<Sentence>> wordToSentence =
        JCasUtil.indexCovering(jCas, WordToken.class, Sentence.class).entrySet().stream()
            .filter(e -> pronounAnnotation.contains(e.getKey()))
            .collect(Collectors.toMap(Entry::getKey, Entry::getValue));
    final Map<Entity, Collection<Sentence>> entityToSentence =
        JCasUtil.indexCovering(jCas, Entity.class, Sentence.class).entrySet().stream()
            .filter(e -> entityAnnotation.contains(e.getKey()))
            .collect(Collectors.toMap(Entry::getKey, Entry::getValue));
    final Map<PhraseChunk, Collection<Sentence>> npToSentence =
        JCasUtil.indexCovering(jCas, PhraseChunk.class, Sentence.class).entrySet().stream()
            .filter(e -> npAnnotation.contains(e.getKey()))
            .collect(Collectors.toMap(Entry::getKey, Entry::getValue));

    // Create a sentence count

    final List<Sentence> sentences = new ArrayList<>(JCasUtil.select(jCas, Sentence.class));
    final Map<Sentence, Integer> sentenceIndex =
        IntStream.range(0, sentences.size())
            .boxed()
            .collect(Collectors.toMap(sentences::get, i -> i));

    // Map mentions to sentence index

    mentions.forEach(
        m -> {
          Collection<Sentence> collection = null;
          switch (m.getType()) {
            case ENTITY:
              collection = entityToSentence.get(m.getAnnotation());
              break;
            case PRONOUN:
              collection = wordToSentence.get(m.getAnnotation());
              break;
            case NP:
              collection = npToSentence.get(m.getAnnotation());
              break;
            default:
              collection = Collections.emptyList();
          }

          setSentence(m, collection, sentenceIndex);
        });
  }

  private void setSentence(
      Mention m, Collection<Sentence> sentenceCollection, Map<Sentence, Integer> sentenceIndex) {
    Sentence s = null;

    if (sentenceCollection != null && !sentenceCollection.isEmpty()) {
      s = sentenceCollection.iterator().next();
    }

    if (s == null) {
      m.setSentence(null);
      m.setSentenceIndex(Integer.MIN_VALUE);
    } else {
      m.setSentence(s);
      m.setSentenceIndex(sentenceIndex.getOrDefault(s, Integer.MIN_VALUE));
    }
  }
}
