content: package org.allenai.scienceparse;
  
  import com.google.common.cache.CacheBuilder;
  import com.google.common.cache.CacheLoader;
  import com.google.common.cache.LoadingCache;
  import com.gs.collections.api.tuple.Pair;
  import com.gs.collections.impl.tuple.Tuples;
  import lombok.extern.slf4j.Slf4j;
  import org.allenai.ml.sequences.crf.CRFModel;
  import org.allenai.scienceparse.ExtractReferences.BibRecordParser;
  import org.allenai.scienceparse.ExtractedMetadata.LabelSpan;
  
  import java.io.*;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.List;
  import java.util.concurrent.TimeUnit;
  import java.util.function.Function;
  import java.util.regex.Matcher;
  import java.util.regex.Pattern;
  
  @Slf4j
  public class CRFBibRecordParser implements BibRecordParser {
    private CRFModel<String, String, String> model;
    public static final String DATA_VERSION = "0.1";
    
    public CRFBibRecordParser(CRFModel<String, String, String> inModel) {
        model = inModel;
    }
    
    public static List<Pair<String, String>> getLabeledLineUMass(String s) {
      final String [] sourceTags = new String [] {"address", "authors", "booktitle", "editor", "institution",
          "journal", "pages", "publisher", "tech", "thesis", "title", "volume", "year"}; //MUST BE SORTED
      final String [] destTags = new String [] {"address", ExtractedMetadata.authorTag, ExtractedMetadata.venueTag,
          "editor", "institution", ExtractedMetadata.venueTag, "pages", "publisher", "tech", ExtractedMetadata.venueTag, 
          ExtractedMetadata.titleTag, "volume", ExtractedMetadata.yearTag};
  
      //pull the ref marker
      s.replaceAll("<ref-marker>.*</ref-marker>", "");
      
      return labelAccordingToTags(s, sourceTags, destTags);
    }
    
    public static List<Pair<String, String>> labelAccordingToTags(
            String s,
            String[] sourceTags,
            String[] destTags
    ) {
      List<String> toks = tokenize(s);
      List<Pair<String, String>> out = new ArrayList<>();
      out.add(Tuples.pair("<S>", "<S>"));
      boolean atStart = false;
      int curTagIndex = -1; //the current source tag that we're in, and labeling for
      for(int i=0; i<toks.size(); i++) {
        String sTok = toks.get(i);
        if(sTok.endsWith(">")) {
          String t = sTok.replaceAll("<", "").replaceAll("/", "").replaceAll(">", "");
          int idx = Arrays.binarySearch(sourceTags, t);
          if(idx >= 0) { //ignore unescaped XML chars in bib
            if(sTok.startsWith("</"))
              curTagIndex = -1; //our tag closed
            else {
              curTagIndex = idx;
              atStart = true;
            }
          }
        } else { //write it out with proper tag
          String tag = "O";
          if(curTagIndex >= 0) {
            tag = destTags[curTagIndex];
            if(i<toks.size()-1) {
              if(toks.get(i+1).equals("</" + sourceTags[curTagIndex] + ">")) { //our tag ends on next word
                if(atStart) { //single-word span
                  tag = "W_" + tag;
                } else {
                  tag = "E_" + tag;
                }
              } else if(atStart) {
                tag = "B_" + tag;
              } else {
                tag = "I_" + tag;
              }
            }
          }
          out.add(Tuples.pair(sTok, tag));
          atStart = false;
        }
      }
      
      out.add(Tuples.pair("</S>", "</S>"));
      return out;
    }
    
    public static List<Pair<String, String>> getLabeledLineCora(String s) {
      final String [] sourceTags = new String [] {"author", "booktitle", "date", "editor", "institution",
          "journal", "location", "note", "pages", "publisher", "tech", "title", "volume"}; //MUST BE SORTED
      final String [] destTags = new String [] {ExtractedMetadata.authorTag, 
          ExtractedMetadata.venueTag, ExtractedMetadata.yearTag, "editor", "institution", 
          ExtractedMetadata.venueTag, "location", "note", "pages", "publisher", ExtractedMetadata.venueTag,
          ExtractedMetadata.titleTag, "volume"};
      
      return labelAccordingToTags(s, sourceTags, destTags);
    }
    
    public static List<Pair<String, String>> getLabeledLineKermit(String s) {
      final String [] sourceTags = new String [] {"author", "booktitle", "date", "editor", 
          "note", "pages", "publisher", "pubPlace", "title", "volume"}; //MUST BE SORTED
      final String [] destTags = new String [] {ExtractedMetadata.authorTag, 
          ExtractedMetadata.venueTag, ExtractedMetadata.yearTag, "editor", "note", "pages",  
          "publisher", "location", ExtractedMetadata.titleTag, "volume"};
      
      s = s.replaceAll("<title level=\"j\"> ([^<]+) </title>", "<booktitle> $1 </booktitle>");
      s = s.replaceAll("<title level=\"m\"> ([^<]+) </title>", "<booktitle> $1 </booktitle>");
      s = s.replaceAll("<title level=\"a\"> ([^<]+) </title>", "<title> $1 </title>");
      s = s.replaceAll("biblScope type=\"vol\" ([^<]+) </biblScope>", "<vol> $1 </vol>");
      s = s.replaceAll("biblScope type=\"pp\" ([^<]+) </biblScope>", "<pages> $1 </pages>");
      s = s.replaceAll("biblScope type=\"issue\" ([^<]+) </biblScope>", "<issue> $1 </issue>");
      
      return labelAccordingToTags(s, sourceTags, destTags);
    }
    
    public static List<List<Pair<String, String>>> labelFromCoraFile(File trainFile) throws IOException {
      return labelFromFile(trainFile, s->getLabeledLineCora(s));
    }
    
    public static List<List<Pair<String, String>>> labelFromUMassFile(File trainFile) throws IOException {
      return labelFromFile(trainFile, s-> getLabeledLineUMass(s));
    }
    
    public static List<List<Pair<String, String>>> labelFromKermitFile(File trainFile) throws IOException {
      return labelFromFile(trainFile, s-> getLabeledLineKermit(s));
    }
    
    
    public static List<List<Pair<String, String>>> labelFromFile(File trainFile, Function<String, List<Pair<String, String>>> gll) throws IOException {
      List<List<Pair<String, String>>> out = new ArrayList<>();
      BufferedReader brIn = new BufferedReader(new InputStreamReader(new FileInputStream(trainFile), "UTF-8"));
      String sLine;
      while((sLine = brIn.readLine())!=null) {
        List<Pair<String,String>> labeledLine = gll.apply(sLine);
        out.add(labeledLine);
      }
      brIn.close();
      return out;
    }
    
    private static List<String> tokenize(String line) { //tokenizes string for use in bib parsing
      //line = line.replaceAll("([a-z0-9])(,|\\.)", "$1 $2"); //break on lowercase alpha or num before . or ,
      String [] toks = line.split(" "); //otherwise not fancy
      return Arrays.asList(toks);
    }
  
    private final LoadingCache<ArrayList<String>, List<String>> bestGuessCache =
        CacheBuilder.newBuilder().
        maximumSize(10240).
        expireAfterAccess(10, TimeUnit.MINUTES).build(
          new CacheLoader<ArrayList<String>, List<String>>() {
            @Override
            public List<String> load(ArrayList<String> key) throws Exception {
              return model.bestGuess(key);
            }
          }
    );
  
    public BibRecord parseRecord(String line) {
      if(Thread.interrupted())
        throw new Parser.ParsingTimeout();
  
      line = line.trim();
      if(line.isEmpty() || line.length() > 2000)
        return null;
      Matcher m = RegexWithTimeout.matcher(ExtractReferences.pBracket, line);
      String citeRegEx = null;
      String shortCiteRegEx = null;
      if(m.matches()) {
        citeRegEx = m.group(1);
        shortCiteRegEx = citeRegEx;
        line = m.group(2);
      } else {
        m = RegexWithTimeout.matcher(ExtractReferences.pDot, line);
        if(m.matches()) {
          citeRegEx = m.group(1);
          shortCiteRegEx = citeRegEx;
          line = m.group(2);
        }
      }
      line = line.trim();
      if(line.isEmpty())
        return null;
      ArrayList<String> toks = new ArrayList<String>();
      toks.add("<S>");
      toks.addAll(tokenize(line));
      toks.add("</S>");
      List<String> labels;
      try{
        labels = bestGuessCache.get(toks);
      } catch(final Exception e) {
        return null;
      }
      labels = PDFToCRFInput.padTagSequence(labels);
      List<LabelSpan> lss = ExtractedMetadata.getSpans(labels);
      
      String title = null;
      String author = null;
      String venue = null;
      String year = null;
      for (LabelSpan ls : lss) {
        //TODO: figure out how to make this code less redundant
        if (title == null && ls.tag.equals(ExtractedMetadata.titleTag)) { //only take the first
          title = PDFToCRFInput.stringAtForStringList(toks, ls.loc);
        } else if (author == null && ls.tag.equals(ExtractedMetadata.authorTag)) {
          author = PDFToCRFInput.stringAtForStringList(toks, ls.loc);
        } else if (venue == null && ls.tag.equals(ExtractedMetadata.venueTag)) {
          venue = PDFToCRFInput.stringAtForStringList(toks, ls.loc);
        } else if (year == null && ls.tag.equals(ExtractedMetadata.yearTag)) {
          year = PDFToCRFInput.stringAtForStringList(toks, ls.loc);
        }
      }
      
      List<String> authors =
              author == null ?
                      null :
                      ExtractReferences.authorStringToList(author);
  //    log.info("authors first extracted: " + ((authors==null)?"":authors.toString()));
  //    log.info("year first extracted: " + year);
      
      int iYear = -1;
      if(year == null && venue != null) { 
        // check venue first, as sometimes year is at end
        Matcher mY = RegexWithTimeout.matcher(ReferencesPredicateExtractor.yearPattern, venue);
        while (mY.find())
          year = mY.group(1);
      }
      if(year == null) {
        // check whole line
        Matcher mY = RegexWithTimeout.matcher(ReferencesPredicateExtractor.yearPattern, line);
        while(mY.find())
          year = mY.group(1);
      }
      if(year != null)
       iYear = ExtractReferences.extractRefYear(year);
  
      if(citeRegEx == null && year != null) {
        shortCiteRegEx = ExtractReferences.getCiteAuthorFromAuthors(authors);
        citeRegEx = shortCiteRegEx + ",? " + ((iYear > 0)?Pattern.quote(iYear + ""):"");
      }
      if(citeRegEx == null || shortCiteRegEx == null || title == null || authors == null || year == null)
        return null;
  //    log.info("cite string " + citeRegEx);
      BibRecord brOut = null;
      try {
        brOut = new BibRecord(
          CRFBibRecordParser.cleanTitle(title),
          Parser.trimAuthors(authors),
          Parser.cleanTitle(venue),
          Pattern.compile(citeRegEx),
          Pattern.compile(shortCiteRegEx),
          iYear);
      } catch (final NumberFormatException e) {
        return null;
      }
      if(iYear==0) //heuristic -- if we don't find year, almost certainly we didn't extract correctly.
        return null;
  //    log.info("returning " + brOut.toString());
      return brOut;
    }
    
    public static String cleanTitle(String title) {
      return title.replaceAll("^\\p{Pi}", "").
          replaceAll("(\\p{Pe}|,|\\.)$", "");
    }
  }
  