package com.pengyifan.pubtator.utils;

import com.google.common.base.Joiner;
import com.pengyifan.brat.BratDocument;
import com.pengyifan.brat.BratEntity;
import com.pengyifan.brat.BratNote;
import com.pengyifan.pubtator.PubTatorDocument;
import com.pengyifan.pubtator.PubTatorMentionAnnotation;

import static com.google.common.base.Preconditions.checkArgument;

public class PubTator2Brat {
  public BratDocument convert(PubTatorDocument pubTatorDocument) {
    BratDocument bratDocument = new BratDocument();
    bratDocument.setDocId(pubTatorDocument.getId());
    bratDocument.setText(Joiner.on("\n").join(
        pubTatorDocument.getTitle(), pubTatorDocument.getAbstract()));
    // mention
    int entityId = 0;
    for(PubTatorMentionAnnotation mention: pubTatorDocument.getMentions()) {
      // test text
      String actual = mention.getText();
      String expected = bratDocument.getText().substring(mention.getStart(), mention.getEnd());
      checkArgument(actual.equals(expected),
          "Not equal. actual[%s %s %s] : expected[%s]", actual, mention.getStart(),
          mention.getEnd(), expected);
      BratEntity entity = new BratEntity();
      entity.setId("T" + entityId++);
      entity.setText(mention.getText());
      entity.setType(mention.getType());
      entity.addSpan(mention.getStart(), mention.getEnd());
      bratDocument.addAnnotation(entity);
      BratNote note = new BratNote();
      note.setId("#" + entity.getId());
      note.setRefId(entity.getId());
      note.setType("ConceptId");
      note.setText(Joiner.on('|').join(mention.getConceptIds()));
      bratDocument.addAnnotation(note);
    }
    return bratDocument;
  }
}

@@ -70,16 +70,95 @@
 Joiner;%0A
+import com.google.common.collect.Lists;%0Aimport com.google.common.collect.Maps;%0A
 import c
@@ -181,32 +181,32 @@
 t.BratDocument;%0A
-
 import com.pengy
@@ -263,16 +263,56 @@
 atNote;%0A
+import com.pengyifan.brat.BratRelation;%0A
 import c
@@ -407,16 +407,120 @@
 otation;
+%0Aimport com.pengyifan.pubtator.PubTatorRelationAnnotation;%0A%0Aimport java.util.List;%0Aimport java.util.Map;
 %0A%0Aimport
@@ -609,16 +609,17 @@
 2Brat %7B%0A
+%0A
   public
@@ -896,24 +896,85 @@
 stract()));%0A
+%0A    Map%3CString, List%3CBratEntity%3E%3E map = Maps.newHashMap();%0A%0A
     // menti
@@ -1875,24 +1875,24 @@
 eptIds()));%0A
-
       bratDo
@@ -1918,16 +1918,888 @@
 n(note);
+%0A%0A      for(String conceptId: mention.getConceptIds()) %7B%0A        if (!map.containsKey(conceptId)) %7B%0A          map.put(conceptId, Lists.newArrayList());%0A        %7D%0A        map.get(conceptId).add(entity);%0A      %7D%0A    %7D%0A    // relation%0A    int relationId = 0;%0A    for(PubTatorRelationAnnotation relation: pubTatorDocument.getRelations()) %7B%0A      List%3CBratEntity%3E e1s = map.get(relation.getConceptId1());%0A      List%3CBratEntity%3E e2s = map.get(relation.getConceptId2());%0A      for(BratEntity e1: e1s) %7B%0A        for (BratEntity e2: e2s) %7B%0A          BratRelation bratRelation = new BratRelation();%0A          bratRelation.setId(%22R%22 + relationId++);%0A          bratRelation.setType(relation.getType());%0A          bratRelation.putArgument(%22Arg1%22, e1.getId());%0A          bratRelation.putArgument(%22Arg2%22, e2.getId());%0A          bratDocument.addAnnotation(bratRelation);%0A        %7D%0A      %7D
 %0A    %7D%0A 

add pubtator to brat relation convert