package org.allenai.scienceparse;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import com.gs.collections.api.list.ImmutableList;
import com.gs.collections.api.map.primitive.DoubleIntMap;
import com.gs.collections.api.map.primitive.MutableDoubleIntMap;
import com.gs.collections.api.map.primitive.MutableFloatIntMap;
import com.gs.collections.api.map.primitive.MutableObjectIntMap;
import com.gs.collections.api.set.primitive.DoubleSet;
import com.gs.collections.api.set.primitive.MutableDoubleSet;
import com.gs.collections.api.tuple.Pair;
import com.gs.collections.api.tuple.primitive.DoubleIntPair;
import com.gs.collections.api.tuple.primitive.FloatIntPair;
import com.gs.collections.impl.factory.primitive.DoubleIntMaps;
import com.gs.collections.impl.factory.primitive.DoubleSets;
import com.gs.collections.impl.factory.primitive.FloatIntMaps;
import com.gs.collections.impl.map.mutable.primitive.ObjectIntHashMap;
import com.gs.collections.impl.tuple.Tuples;
import org.allenai.scienceparse.pdfapi.PDFDoc;
import org.allenai.scienceparse.pdfapi.PDFLine;
import org.allenai.scienceparse.pdfapi.PDFPage;

import lombok.extern.slf4j.Slf4j;
import org.allenai.scienceparse.pdfapi.PDFToken;

@Slf4j
public class PDFDocToPartitionedText {
  /**
   * Returns list of strings representation of this file.  Breaks new lines when pdf line break larger than threshold.
   * All original line breaks indicated by <lb>
   *
   * @param pdf
   * @return
   */
  public static List<String> getRaw(PDFDoc pdf) {
    ArrayList<String> out = new ArrayList<>();

    StringBuilder s = new StringBuilder();
    PDFLine prevLine = null;
    double qLineBreak = getRawBlockLineBreak(pdf);
    for (PDFPage p : pdf.getPages()) {
      for (PDFLine l : p.getLines()) {
        if (breakSize(l, prevLine) > qLineBreak) {
          String sAdd = s.toString();
          if (sAdd.endsWith("<lb>"))
            sAdd = sAdd.substring(0, sAdd.length() - 4);
          out.add(sAdd);
          s = new StringBuilder();
        }
        String sAdd = lineToString(l);
        if (sAdd.length() > 0) {
          s.append(sAdd);
          s.append("<lb>");
        }
        prevLine = l;
      }
      //HACK(dcdowney): always break on new page.  Should be safe barring "bad breaks" I think
      if (s.length() > 0) {
        String sAdd = s.toString();
        if (sAdd.endsWith("<lb>"))
          sAdd = sAdd.substring(0, sAdd.length() - 4);
        out.add(sAdd);
        s = new StringBuilder();
      }
    }
    return out;
  }

  public static float breakSize(PDFLine l2, PDFLine l1) {
    if (l2 == null || l1 == null)
      return 0.0f;
    float h1 = PDFToCRFInput.getH(l1);
    float h2 = PDFToCRFInput.getH(l2);
    return (PDFToCRFInput.getY(l2, true) - PDFToCRFInput.getY(l1, false)) / Math.min(h1, h2);
  }
  
  private static List<Double> getBreaks(PDFPage p) {
    PDFLine prevLine = null;
    ArrayList<Double> breaks = new ArrayList<>();
    for (PDFLine l : p.getLines()) {
      double bs = breakSize(l, prevLine);
      if (bs > 0) { //<= 0 due to math, tables, new pages, should be ignored
        breaks.add(bs);
      }
      prevLine = l;
    }
    breaks.sort(Double::compare);
    return breaks;
  }
  
  private static List<Double> getBreaks(PDFDoc pdf) {
    ArrayList<Double> breaks = new ArrayList<>();
    for (PDFPage p : pdf.getPages()) {
      breaks.addAll(getBreaks(p));
    }
    breaks.sort(Double::compare);
    return breaks;
  }
  
  public static double getReferenceLineBreak(PDFDoc pdf) {
    List<Double> breaks = getBreaks(pdf);
    if(breaks.isEmpty())
      return 1.0;
    int idx = (7 * breaks.size()) / 9; //hand-tuned threshold good for breaking references
    return breaks.get(idx);
  }
  
  public static double getRawBlockLineBreak(PDFDoc pdf) {
    List<Double> breaks = getBreaks(pdf);
    if(breaks.isEmpty())
      return 1.0;
    int idx = (7 * breaks.size()) / 9; //hand-tuned threshold good for breaking papers
    return breaks.get(idx);
  }
  
  public static double getFirstPagePartitionBreak(PDFPage pdf) {
    List<Double> breaks = getBreaks(pdf);
    if(breaks.isEmpty())
      return 1.0;
    int idx = (3 * breaks.size()) / 6; //hand-tuned threshold good for breaking first pages (abstracts)
    return breaks.get(idx) + 0.50;
  }

  private static String lineToString(PDFLine l) {
    StringBuilder sb = new StringBuilder();
    for(PDFToken token : l.tokens) {
      sb.append(token.token);
      sb.append(' ');
    }
    return sb.toString().trim();
  }

  private static String cleanLine(String s) {
    s = s.replaceAll("\r|\t|\n", " ").trim();
    while (s.contains("  "))
      s = s.replaceAll("  ", " ");
    return s;
  }

  public static String getFirstTextBlock(PDFDoc pdf) {
    PDFPage fp = pdf.pages.get(0);
    double fpp = getFirstPagePartitionBreak(fp);
    StringBuilder out = new StringBuilder();
    PDFLine prevLine = null;
    boolean first = true;
    for(PDFLine l : fp.lines) {
      if(first) {
        first=false; //skip the first line (heuristic)
        continue;
      }
      if (breakSize(l, prevLine) > fpp) {
        if(out.length() > 400) { //hand-tuned threshold of min abstract length
          return out.toString().trim();
        } else {
          out.delete(0, out.length());
          out.append(' ');
          out.append(cleanLine(lineToString(l)));
        }
      }
      else {
        out.append(' ');
        out.append(cleanLine(lineToString(l)));
      }
      prevLine = l;
    }
    return "";
  }


  private final static Pattern inLineAbstractPattern =
    Pattern.compile("^abstract ?\\p{P}?", Pattern.UNICODE_CASE | Pattern.CASE_INSENSITIVE);

  private final static Pattern[] generalAbstractCleaners = new Pattern[] {
    Pattern.compile("Key ?words(:| |\\.).*$", Pattern.UNICODE_CASE),
    Pattern.compile("KEY ?WORDS(:| |\\.).*$", Pattern.UNICODE_CASE),
    Pattern.compile("Key ?Words(:| |\\.).*$", Pattern.UNICODE_CASE),
    Pattern.compile("(1|I)\\.? Introduction.*$", Pattern.UNICODE_CASE | Pattern.CASE_INSENSITIVE),
    Pattern.compile("Categories and Subject Descriptors.*$", Pattern.UNICODE_CASE),
    Pattern.compile("0 [1-2][0-9]{3}.*$", Pattern.UNICODE_CASE),
    Pattern.compile("Contents.*$", Pattern.UNICODE_CASE),
    Pattern.compile("Index terms\\p{P}.*$", Pattern.UNICODE_CASE | Pattern.CASE_INSENSITIVE),
  };
  private final static Pattern paragraphAbstractCleaner =
    Pattern.compile("^summary ?\\p{P}?", Pattern.UNICODE_CASE | Pattern.CASE_INSENSITIVE);

  public static String getAbstract(List<String> raw, PDFDoc pdf) {
    boolean inAbstract = false;
    StringBuilder out = new StringBuilder();
    for(String s : raw) {
      if(inAbstract) {
        if(s.length() < 20)
          break;
        else {
          out.append(' ');
          out.append(s.trim());
        }
      }

      if(s.toLowerCase().contains("abstract") && s.length() < 10) {
        inAbstract = true;
      } else if(s.toLowerCase().contains("a b s t r a c t")) {
        inAbstract = true;
      } else if(RegexWithTimeout.matcher(inLineAbstractPattern, s).find()) {
        out.append(RegexWithTimeout.matcher(inLineAbstractPattern, s).replaceFirst(""));
        inAbstract = true;
      }
    }
    String abs = out.toString().trim();
    if(abs.length()==0) {
      //we didn't find an abstract.  Pull out the first paragraph-looking thing.
      abs = getFirstTextBlock(pdf);
      abs = RegexWithTimeout.matcher(paragraphAbstractCleaner, abs).replaceFirst("");
    }
    
    // remove keywords, intro from abstract
    for(Pattern p : generalAbstractCleaners) {
      abs = RegexWithTimeout.matcher(p, abs).replaceFirst("");
    }

    abs = abs.replaceAll("- ", "");
    return abs;
  }
  
  private static boolean lenientRefStart(PDFLine l, PDFLine prevLine, double qLineBreak) {
    final PDFToken firstToken = l.tokens.get(0);

    return (
          firstToken.token.equals("[1]") ||
          firstToken.token.equals("1.")
        ) && l.tokens.size() > 1 && (
            PDFToCRFInput.getX(l.tokens.get(1), true) > firstToken.fontMetrics.spaceWidth ||
            breakSize(l, prevLine) > qLineBreak
        );
  }
  
  public static Set<String> referenceHeaders = new HashSet<String>(Arrays.asList(
      "references",
      "citations",
      "bibliography",
      "reference",
      "bibliographie"));

  private static Pattern referenceStartPattern =
      Pattern.compile("^\\d{1,2}(\\.|\\s|$)|^\\[.+?\\]");

  private static boolean gapAcrossMiddle(PDFToken t1, PDFToken t2, PDFPage p, float lineSpaceWidth) {
    double gap = PDFToCRFInput.getXGap(t1, t2);
    double pageCenter = p.getPageWidth() / 2.0;
    double gapCenter = (PDFToCRFInput.getX(t1, false) + PDFToCRFInput.getX(t2, true))/2.0;
    return gap > 5*lineSpaceWidth &&
            Math.abs(gapCenter - pageCenter) < 50*lineSpaceWidth; //lenient on center since margins might differ
  }

  /**
   * The lower-level processing sometimes fails to detect column breaks and/or fails to order
   * to column lines correctly. This function attempts to repair that, returning column-broken lines
   * ordered left-to-right, top-to-bottom.
   * @param lines
   * @return
   */
  public static List<Pair<PDFPage, PDFLine>> repairColumns(List<Pair<PDFPage, PDFLine>> lines) {
    List<Pair<PDFPage, PDFLine>> out = new ArrayList<>();
    List<PDFLine> linesA = new ArrayList<>(); //holds new, broken lines
    PDFPage prevPage = null;

    ArrayList<Double> gaps = new ArrayList<>();

    double meanSpaceWidth = 0.0;

    for(final Pair<PDFPage, PDFLine> pageLinePair : lines) {
      PDFLine line = pageLinePair.getTwo();
      PDFPage page = pageLinePair.getOne();
      if(page != prevPage && prevPage != null) {
        final PDFPage comparePage = prevPage;
        linesA.sort((line1, line2) -> Double.compare(lineSorter(line1, comparePage), lineSorter(line2, comparePage)));
        for(PDFLine linea : linesA)
          out.add(Tuples.pair(prevPage, linea));
        linesA = new ArrayList<>();
      }
      List<PDFToken> lineAcc = new ArrayList<>();
      PDFToken prevToken = null;
      for(final PDFToken token : line.tokens) {
        if(prevToken != null) {
          if(gapAcrossMiddle(prevToken, token, page, prevToken.fontMetrics.spaceWidth)) {
            gaps.add((PDFToCRFInput.getX(prevToken,false) + PDFToCRFInput.getX(token, true))/2.0);
            meanSpaceWidth += prevToken.fontMetrics.spaceWidth;
            linesA.add(PDFLine.builder().tokens(new ArrayList<>(lineAcc)).build());
            lineAcc = new ArrayList<>();
          }
        }
        lineAcc.add(token);
        prevToken = token;
      }
      if(lineAcc.size() > 0)
        linesA.add(PDFLine.builder().tokens(new ArrayList<>(lineAcc)).build());
      prevPage = page;
    }
    //if gaps are consistent and there are at least ten of them:
    boolean useReorderedLines = false;
    if(gaps.size() > 10) {
      useReorderedLines = true;
      double mean = 0.0;
      meanSpaceWidth /= gaps.size();
      for(double d : gaps) {
        mean += d;
      }
      mean /= gaps.size();
      for(double d : gaps) {
        if(Math.abs(d - mean) > 3*meanSpaceWidth) {
          useReorderedLines = false;
          break;
        }
      }
    }
    if(useReorderedLines) {
      final PDFPage comparePage = prevPage;
      linesA.sort((line1, line2) -> Double.compare(lineSorter(line1, comparePage), lineSorter(line2, comparePage)));
      for (PDFLine linea : linesA)
        out.add(Tuples.pair(prevPage, linea));
      log.info("using re-ordered lines: " + out.size());
      return out;
    }
    else
      return lines;
  }

  private static double lineSorter(PDFLine line, PDFPage p) {
    return 1E8*(firstCol(line, p)?0.0:1.0) + PDFToCRFInput.getY(line, true);
  }

  private static boolean firstCol(PDFLine line, PDFPage p) {
    double pageThird = p.getPageWidth() / 3.0;
    return PDFToCRFInput.getX(line, true) < pageThird;
  }

  private static boolean referenceIsSplit(String lineOne, String lineTwo) {
    Matcher lineOneMatcher = RegexWithTimeout.matcher(referenceStartPattern, lineOne);
    return (lineOneMatcher.find() &&
            lineOneMatcher.end() == lineOne.length() &&
            !RegexWithTimeout.matcher(referenceStartPattern, lineTwo).find());
  }

  /**
   * Returns best guess of list of strings representation of the references of this file,
   * intended to be one reference per list element, using spacing and indentation as cues
   */
  public static List<String> getRawReferences(PDFDoc pdf) {
    PDFLine prevLine = null;
    boolean inRefs = false;
    boolean foundRefs = false;
    double qLineBreak = getReferenceLineBreak(pdf);
    boolean lenient = false;

    // Find reference lines in the document
    List<Pair<PDFPage, PDFLine>> referenceLines = new ArrayList<>();
    int totalLines = 0;

    for(int pass=0;pass<2;pass++) {
      int passLines = 0;
      if(pass==1)
        if(foundRefs)
          break;
        else
          lenient=true; //try harder this time.
      for (PDFPage p : pdf.getPages()) {
        double farLeft = Double.MAX_VALUE; //of current column
        double farRight = -1.0; //of current column
        for (PDFLine l : p.getLines()) {
          if (!inRefs && (l != null && l.tokens != null && l.tokens.size() > 0)) {
            if (
              l.tokens.get(l.tokens.size() - 1).token != null &&
              referenceHeaders.contains(
                  l.tokens.get(l.tokens.size() - 1).token.trim().toLowerCase().replaceAll("\\p{Punct}*$", "")) &&
              l.tokens.size() < 5
            ) {
              inRefs = true;
              foundRefs = true;
              prevLine = l;
              continue; //skip this line
            }
            else if (lenient && passLines > totalLines / 4) { //used if we don't find refs on first pass; must not be in first 1/4 of doc
              if (lenientRefStart(l, prevLine, qLineBreak)) {
                inRefs = true;
                foundRefs = true;
                //DON'T skip this line.
              }
            }
          }
          if (inRefs)
            referenceLines.add(Tuples.pair(p, l));
          prevLine = l;
        }
        if(pass==0)
          totalLines++;
        passLines++;
      }
    }

    referenceLines = repairColumns(referenceLines);

    // split reference lines into columns
    final List<List<PDFLine>> referenceLinesInColumns = new ArrayList<>();
    PDFPage lastPage = null;
    double currentColumnBottom = Double.MAX_VALUE;
    for(final Pair<PDFPage, PDFLine> pageLinePair : referenceLines) {
      final PDFPage p = pageLinePair.getOne();
      final PDFLine l = pageLinePair.getTwo();

      // remove empty lines
      if(l.tokens.isEmpty())
        continue;

      // Cut into columns. One column is a set of lines with continuously increasing Y coordinates.
      final double lineTop = PDFToCRFInput.getY(l, true);
      final double lineBottom = PDFToCRFInput.getY(l, false);
      if (p != lastPage || lineTop < currentColumnBottom) {
        final List<PDFLine> newColumn = new ArrayList<>();
        newColumn.add(l);
        referenceLinesInColumns.add(newColumn);
        currentColumnBottom = lineBottom;
      } else {
        referenceLinesInColumns.get(referenceLinesInColumns.size() - 1).add(l);
        currentColumnBottom = lineBottom;
      }
      lastPage = p;
    }

    // parse each column into output
    // We assume that the indentation of the first line of every column marks the start of a
    // reference (unless that indentation happens only once)
    final List<String> out = new ArrayList<String>();
    for(final List<PDFLine> column : referenceLinesInColumns) {
      // find indentation levels
      final MutableFloatIntMap left2count = FloatIntMaps.mutable.empty();
      for(final PDFLine l : column) {
        final float left = PDFToCRFInput.getX(l, true);
        final float lineSpaceWidth = l.tokens.get(0).fontMetrics.spaceWidth;

        // find an indentation level that's close to this one
        // This is not proper clustering, but I don't think we need it.
        float foundLeft = left;
        for(final float indentLevel : left2count.keySet().toArray()) {
          if(Math.abs(indentLevel - left) < lineSpaceWidth) {
            foundLeft = indentLevel;
            break;
          }
        }

        final int oldCount = left2count.getIfAbsent(foundLeft, 0);
        left2count.remove(foundLeft);
        left2count.put(
            (foundLeft * oldCount + left) / (oldCount + 1),
            oldCount + 1);
      }

      // find the indentation that starts a reference
      float startReferenceIndentation = -1000; // default is some number that's definitely more than one space width away from a realistic indent
      final ImmutableList<FloatIntPair> startReferenceIndentCandidates =
          left2count.keyValuesView().toSortedListBy(pair -> -pair.getTwo()).take(2).toImmutable();
      if(startReferenceIndentCandidates.size() > 1) {
        final FloatIntPair first = startReferenceIndentCandidates.get(0);
        final FloatIntPair second = startReferenceIndentCandidates.get(1);

        // find lines that look like they start references, and use them to determine which indent
        // starts a reference
        float firstIndentSeen = -1;
        for(final PDFLine l : column) {
          float left = PDFToCRFInput.getX(l, true);
          final float lineSpaceWidth = l.tokens.get(0).fontMetrics.spaceWidth;

          // snap to candidate indentations
          if(Math.abs(first.getOne() - left) < lineSpaceWidth)
            left = first.getOne();
          else if(Math.abs(second.getOne() - left) < lineSpaceWidth)
            left = second.getOne();
          else
            continue; // only consider candidates

          if(firstIndentSeen < 0)
            firstIndentSeen = left;

          final String lineAsString = lineToString(l);
          if(referenceStartPattern.matcher(lineAsString).find()) {
            startReferenceIndentation = left;
            break;
          }
        }

        if(startReferenceIndentation < 0) {
          startReferenceIndentation = second.getOne();  // pick the one that's less common
        }
      }

      // find vertical spaces that might separate references
      final float breakSizeTolerance = 0.1f;
      MutableFloatIntMap vspace2count = FloatIntMaps.mutable.empty();
      for(int i = 1; i < column.size(); ++i) {
        final PDFLine l1 = column.get(i - 1);
        final PDFLine l2 = column.get(i);
        final float thisVspace = breakSize(l2, l1); // vspace is measured in line heights

        // if it's close to another break size, we should cluster them
        float foundVspace = thisVspace;
        for(final float v : vspace2count.keySet().toArray()) {
          if(Math.abs(v - thisVspace) <= breakSizeTolerance) {
            foundVspace = v;
            break;
          }
        }

        final int oldCount = vspace2count.getIfAbsent(foundVspace, 0);
        vspace2count.remove(foundVspace);
        vspace2count.put(
            (foundVspace * oldCount + thisVspace) / (oldCount + 1),
            oldCount + 1);
      }
      // filter down to reasonable vspaces
      final long numberOfBreaks = vspace2count.sum();
      vspace2count = vspace2count.select((vspace, count) ->
          count >= numberOfBreaks / 5 &&  // must account for 20% of breaks
          count > 1 &&                    // must occur at least once
          vspace < 3                      // more than 3 lines of gap is crazy
      );

      float breakReferenceVspace = -1.0f;
      if(column.size() > 5 && vspace2count.size() >= 2) // only if we have at least 5 lines
        breakReferenceVspace = vspace2count.keySet().max();

      int linesGrouped = 0; // We never group more than six lines together at a time.
      final StringBuilder builder = new StringBuilder();
      prevLine = null;
      for(final PDFLine l : column) {
        final float left = PDFToCRFInput.getX(l, true);
        final String lineAsString = lineToString(l);
        final float lineSpaceWidth = l.tokens.get(0).fontMetrics.spaceWidth;
        linesGrouped += 1;

        final boolean br =
          linesGrouped >= 6 || (
            breakReferenceVspace < 0 && // if we have vspace, we don't use indent
            startReferenceIndentation > 0 &&
            Math.abs(left-startReferenceIndentation) < lineSpaceWidth
          ) || referenceStartPattern.matcher(lineAsString).find() || (
            prevLine != null &&
            breakReferenceVspace > 0 &&
            breakSize(l, prevLine) >= breakReferenceVspace - breakSizeTolerance
          );

        if(br) {
          // save old line
          final String outLine = cleanLine(builder.toString());
          if(!outLine.isEmpty())
            out.add(outLine);
          // start new line
          builder.setLength(0);
          builder.append(lineAsString);
          linesGrouped = 1;
        } else {
          if (!builder.toString().isEmpty())
            builder.append("<lb>");
          builder.append(lineAsString);
        }

        prevLine = l;
      }
      // save last line
      final String outLine = cleanLine(builder.toString());
      if(!outLine.isEmpty())
        out.add(outLine);
    }

    // If two columns were found incorrectly, the out array may consist of alternating
    // reference numbers and reference information. In this case, combine numbers with
    // the content that follows.
    List<String> mergedRefs = new ArrayList<String>();
    int i=0;
    while (i<out.size()) {
      String thisRef = out.get(i);
      String nextRef = i < out.size()-1 ? out.get(i+1) : null;
      if (nextRef != null && referenceIsSplit(thisRef, nextRef)) {
        mergedRefs.add(String.join(" ", thisRef, nextRef));
        i += 2;
      } else {
        mergedRefs.add(thisRef);
        i += 1;
      }
    }

    return mergedRefs;
  }
}
