// Dstl (c) Crown Copyright 2017
package uk.gov.dstl.baleen.consumers.csv.internals;

import java.util.List;
import java.util.Map;

import org.apache.uima.UimaContext;
import org.apache.uima.fit.util.JCasUtil;
import org.apache.uima.jcas.JCas;
import org.apache.uima.resource.ResourceInitializationException;

import uk.gov.dstl.baleen.consumers.csv.AbstractCsvConsumer;
import uk.gov.dstl.baleen.types.language.Sentence;
import uk.gov.dstl.baleen.types.semantic.Relation;

/**
 * Write relations to CSV.
 *
 * <p>Format is:
 *
 * <ul>
 *   <li>source
 *   <li>sentence
 *   <li>type
 *   <li>subtype
 *   <li>source
 *   <li>target
 *   <li>sourceText
 *   <li>targetText
 *   <li>sourceType
 *   <li>targetType
 *   <li>confidence
 * </ul>
 *
 * @baleen.javadoc
 */
public class CsvRelation extends AbstractCsvConsumer {

  @Override
  public void doInitialize(UimaContext aContext) throws ResourceInitializationException {

    super.doInitialize(aContext);

    write(
        "source",
        "sentence",
        "type",
        "subtype",
        "source",
        "target",
        "sourceText",
        "targetText",
        "sourceType",
        "targetType",
        "confidence");
  }

  @Override
  protected void write(JCas jCas) {
    final String source = getDocumentAnnotation(jCas).getSourceUri();

    final Map<Relation, List<Sentence>> coveringSentence =
        JCasUtil.indexCovering(jCas, Relation.class, Sentence.class);

    JCasUtil.select(jCas, Relation.class).stream()
        .map(
            r -> {
              String sentence = "";
              final List<Sentence> sentences = coveringSentence.get(r);
              if (!sentences.isEmpty()) {
                sentence = sentences.iterator().next().getCoveredText();
              }

              return new String[] {
                source,
                sentence,
                r.getRelationshipType(),
                r.getRelationSubType(),
                normalize(r.getSource().getValue()),
                normalize(r.getTarget().getValue()),
                normalize(r.getSource().getCoveredText()),
                normalize(r.getTarget().getCoveredText()),
                r.getSource().getType().getShortName(),
                r.getTarget().getType().getShortName(),
                Double.toString(r.getConfidence())
              };
            })
        .forEach(this::write);
  }
}
