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

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

import java.util.ArrayList;
import java.util.List;

import org.apache.uima.analysis_engine.AnalysisEngine;
import org.apache.uima.analysis_engine.AnalysisEngineProcessException;
import org.apache.uima.fit.factory.ExternalResourceFactory;
import org.apache.uima.fit.util.JCasUtil;
import org.apache.uima.resource.ExternalResourceDescription;
import org.apache.uima.resource.ResourceInitializationException;
import org.junit.Test;

import uk.gov.dstl.baleen.annotators.language.MaltParser;
import uk.gov.dstl.baleen.annotators.language.OpenNLP;
import uk.gov.dstl.baleen.annotators.language.OpenNLPParser;
import uk.gov.dstl.baleen.annotators.testing.AbstractMultiAnnotatorTest;
import uk.gov.dstl.baleen.resources.SharedOpenNLPModel;
import uk.gov.dstl.baleen.types.common.Person;
import uk.gov.dstl.baleen.types.language.Interaction;
import uk.gov.dstl.baleen.types.semantic.Entity;
import uk.gov.dstl.baleen.types.semantic.Location;
import uk.gov.dstl.baleen.types.semantic.Relation;
import uk.gov.dstl.baleen.types.semantic.Temporal;

public class UbmreConstituentRelationshipTest extends AbstractMultiAnnotatorTest {

  @Override
  protected AnalysisEngine[] createAnalysisEngines() throws ResourceInitializationException {
    ExternalResourceDescription parserChunkingDesc =
        ExternalResourceFactory.createNamedResourceDescription(
            "parserChunking", SharedOpenNLPModel.class);

    // Add in the OpenNLP implementation too, as its a prerequisite
    // (in theory we should test OpenNLPParser in isolation, but in practise it as this as a
    // dependency
    // so better test they work together)

    ExternalResourceDescription tokensDesc =
        ExternalResourceFactory.createNamedResourceDescription("tokens", SharedOpenNLPModel.class);
    ExternalResourceDescription sentencesDesc =
        ExternalResourceFactory.createNamedResourceDescription(
            "sentences", SharedOpenNLPModel.class);
    ExternalResourceDescription posDesc =
        ExternalResourceFactory.createNamedResourceDescription("posTags", SharedOpenNLPModel.class);
    ExternalResourceDescription chunksDesc =
        ExternalResourceFactory.createNamedResourceDescription(
            "phraseChunks", SharedOpenNLPModel.class);

    return asArray(
        createAnalysisEngine(
            OpenNLP.class,
            "tokens",
            tokensDesc,
            "sentences",
            sentencesDesc,
            "posTags",
            posDesc,
            "phraseChunks",
            chunksDesc),
        createAnalysisEngine(OpenNLPParser.class, "parserChunking", parserChunkingDesc),
        createAnalysisEngine(MaltParser.class),
        createAnalysisEngine(UbmreConstituent.class));
  }

  @Test
  public void test() throws AnalysisEngineProcessException, ResourceInitializationException {
    // From the above pipeline we'll have WordToken, Setence and Depedency
    // We still need Entity and Interaction

    jCas.setDocumentText("Jon and Chris will visit London on Monday.");

    Person jon = new Person(jCas);
    jon.setBegin(0);
    jon.setEnd(3);
    jon.addToIndexes();

    Person chris = new Person(jCas);
    chris.setBegin(8);
    chris.setEnd(13);
    chris.addToIndexes();

    Interaction visit = new Interaction(jCas);
    visit.setBegin(19);
    visit.setEnd(24);
    visit.setRelationshipType("visit");
    visit.setValue("visit");
    visit.addToIndexes();

    Location london = new Location(jCas);
    london.setBegin(25);
    london.setEnd(31);
    london.addToIndexes();

    Temporal monday = new Temporal(jCas);
    monday.setBegin(35);
    monday.setEnd(41);
    monday.addToIndexes();

    processJCas();

    List<Relation> relations = new ArrayList<>(JCasUtil.select(jCas, Relation.class));

    // jon-visit-london, jon-visit-monday
    // chris-visit-london, chris-visit-monday.

    assertEquals(4, relations.size());

    assertTrue(relations.stream().anyMatch(r -> assertRelation(r, jon, london, "visit")));
    assertTrue(relations.stream().anyMatch(r -> assertRelation(r, jon, monday, "visit")));
    assertTrue(relations.stream().anyMatch(r -> assertRelation(r, chris, monday, "visit")));
    assertTrue(relations.stream().anyMatch(r -> assertRelation(r, london, chris, "visit")));
  }

  private boolean assertRelation(Relation r, Entity a, Entity b, String type) {
    return (r.getSource() == a && r.getTarget() == b || r.getSource() == b && r.getTarget() == a)
        && r.getRelationshipType().equalsIgnoreCase(type);
  }
}
