package org.allenai.scienceparse;

import com.gs.collections.api.map.primitive.ObjectDoubleMap;
import com.gs.collections.api.tuple.Pair;
import com.gs.collections.impl.map.mutable.primitive.ObjectDoubleHashMap;
import com.gs.collections.impl.tuple.Tuples;
import org.allenai.word2vec.Searcher;
import org.allenai.word2vec.Word2VecModel;
import lombok.extern.slf4j.Slf4j;
import lombok.val;
import org.allenai.datastore.Datastore;
import org.allenai.ml.sequences.crf.CRFPredicateExtractor;
import org.allenai.scienceparse.pdfapi.PDFToken;

import java.io.IOException;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

@Slf4j
public class PDFPredicateExtractor implements CRFPredicateExtractor<PaperToken, String> {

  public static final List<String> stopWords = Arrays.asList("a", "an", "the", "in", "of", "for", "from", "and", "as", "but",
    "to");
  public static final HashSet<String> stopHash = new HashSet<String>(stopWords);

  private final ParserLMFeatures lmFeats;
  private final Searcher word2vecSearcher;

  public PDFPredicateExtractor() {
    this(null);
  }

  public PDFPredicateExtractor(ParserLMFeatures plf) {
    try {
      final Path word2VecModelPath =
              Datastore.apply().filePath("org.allenai.scienceparse", "Word2VecModel.bin", 1);
      word2vecSearcher = WordVectorCache.searcherForPath(word2VecModelPath);
    } catch(final IOException e) {
      throw new RuntimeException(e);
    }

    lmFeats = plf;
  }

  public static List<String> getCaseMasks(String tok) {
    Pattern Xxx = Pattern.compile("[A-Z][a-z]*");
    Pattern xxx = Pattern.compile("[a-z]+");
    Pattern letters = Pattern.compile("[a-zA-Z]+");
    Pattern dig = Pattern.compile("[0-9]+");
    Pattern hasNum = Pattern.compile(".*[0-9]+.*");
    Pattern letterDot = Pattern.compile("[A-Z]\\.");
    Pattern hasNonAscii = Pattern.compile(".*[^\\p{ASCII}]+.*");
    Pattern wordColon = Pattern.compile("[a-zA-Z]+:");
    Pattern hasAt = Pattern.compile(".*@.*");

    List<Pattern> pats = Arrays.asList(Xxx, xxx, letters, dig, hasNum, letterDot, hasNonAscii, wordColon, hasAt);
    List<String> feats = Arrays.asList("%Xxx", "%xxx", "%letters", "%dig", "%hasNum", "%letDot", "%hasNonAscii", "%capWordColon", "%hasAt");
    ArrayList<String> out = new ArrayList<String>();
    for (int i = 0; i < pats.size(); i++) {
      Pattern p = pats.get(i);
      if (p.matcher(tok).matches()) {
        out.add(feats.get(i));
      }
    }
    return out;
  }

  public static boolean isStopWord(String tok) {
    return stopHash.contains(tok);
  }

  public static float getY(PaperToken t, boolean upper) {
    if (upper)
      return t.getPdfToken().bounds.get(1);
    else
      return t.getPdfToken().bounds.get(3);
  }

  public static double smoothFreq(String tok, ObjectDoubleHashMap<String> hm) {
    double freq = hm.get(tok);
    if (freq > 0.0)
      freq -= 0.6;
    return Math.log10(freq + 0.1);
  }

  private static float height(PDFToken t) {
    return t.bounds.get(3) - t.bounds.get(1);
  }

  private static float width(PDFToken t) {
    return t.bounds.get(0) - t.bounds.get(2);
  }

  public static float getExtreme(List<PaperToken> toks, TokenPropertySelector s, boolean max) {
    float adj = -1.0f;
    float extremeSoFar = Float.NEGATIVE_INFINITY;
    if (max) {
      adj = 1.0f;
    }
    for (PaperToken pt : toks) {
      float propAdj = s.getProp(pt) * adj;
      if (propAdj > extremeSoFar) {
        extremeSoFar = propAdj;
      }
    }
    return extremeSoFar * adj;
  }

  public static float linearNormalize(float f, Pair<Float, Float> rng) {
    if (Math.abs(rng.getTwo() - rng.getOne()) < 0.00000001)
      return 0.5f;
    else
      return (f - rng.getOne()) / (rng.getTwo() - rng.getOne());
  }

  public static Pair<Float, Float> getExtrema(List<PaperToken> toks, TokenPropertySelector s) {
    Pair<Float, Float> out = Tuples.pair(getExtreme(toks, s, false), getExtreme(toks, s, true));
    return out;
  }

  public static float getFixedFont(PaperToken t) {
    float s = t.getPdfToken().fontMetrics.ptSize;
    if (s > 30.0f) //assume it's an error
      return 11.0f;
    else
      return s;
  }

  public double logYDelt(float y1, float y2) {
    return Math.log(Math.max(y1 - y2, 0.00001f));
  }

  // String.format() ends up taking a very long time at scale, so we pre-compute all the
  // String.format() calls we might need and re-use them.
  public static final String[] wordEmbeddingFeatureNames;
  static {
    wordEmbeddingFeatureNames = new String[1000];
    for (int i = 0; i < wordEmbeddingFeatureNames.length; ++i)
      wordEmbeddingFeatureNames[i] = String.format("%%emb%03d", i);
  }

  //assumes start/stop padded
  @Override
  public List<ObjectDoubleMap<String>> nodePredicates(List<PaperToken> elems) {
    List<ObjectDoubleMap<String>> out = new ArrayList<>();
    Pair<Float, Float> hBounds = getExtrema(elems.subList(1, elems.size() - 1), (PaperToken t) -> {
      return height(t.getPdfToken());
    });
    Pair<Float, Float> fBounds = getExtrema(elems.subList(1, elems.size() - 1), (PaperToken t) -> {
      return getFixedFont(t);
    });

    for (int i = 0; i < elems.size(); i++) {
      ObjectDoubleHashMap<String> m = new ObjectDoubleHashMap<String>();
      float prevFont = -10.0f;
      float nextFont = -10.0f;
      float prevHeight = -10.0f;
      float nextHeight = -10.0f;
      float prevY = 0.0f;
      float nextY = -1000000.0f;

      int prevLine = -1;
      int nextLine = -1;
      if (i == 0)
        m.put("<S>", 1.0);
      else if (i == elems.size() - 1)
        m.put("</S>", 1.0);
      else {
        if (i != 1) {
          prevLine = elems.get(i - 1).getLine();
          prevFont = getFixedFont(elems.get(i - 1));
          prevHeight = height(elems.get(i - 1).getPdfToken());
          prevY = getY(elems.get(i - 1), false);
        }
        if (i != elems.size() - 2) {
          nextLine = elems.get(i + 1).getLine();
          nextFont = getFixedFont(elems.get(i + 1));
          nextHeight = height(elems.get(i + 1).getPdfToken());
          nextY = getY(elems.get(i + 1), true);
        } else {
          nextY = getY(elems.get(i), false) + height(elems.get(i).getPdfToken()); //guess that next line is height units below
        }
        float font = getFixedFont(elems.get(i));
        float h = height(elems.get(i).getPdfToken());
        int line = elems.get(i).getLine();
        //font-change forward (fcf) or backward (fcb):
        if (font != prevFont)
          m.put("%fcb", 1.0); //binary, 1.0 if there is a font change forward, 0.0 otherwise
        if (font != nextFont)
          m.put("%fcf", 1.0); //font change backward
        if (line != prevLine) {
          m.put("%lcb", 1.0); //line change backward
          m.put("%hGapB", logYDelt(getY(elems.get(i), true), prevY)); //height gap backward
        }
        if (line != nextLine) {
          m.put("%lcf", 1.0); //line change forward
          m.put("%hGapF", logYDelt(nextY, getY(elems.get(i), false))); //height gap forward
        }
        if (Math.abs(Math.abs(nextHeight - h) / Math.abs(nextHeight + h)) > 0.1) { //larger than ~20% height change forward
          m.put("%hcf", 1.0);
        }
        if (Math.abs(Math.abs(prevHeight - h) / Math.abs(prevHeight + h)) > 0.1) { //larger than ~20% height change backward
          m.put("%hcb", 1.0);
        }

        //font value:
        float relativeF = linearNormalize(font, fBounds);
        m.put("%font", relativeF); //font size normalized relative to doc

        m.put("%line", Math.min(line, 10.0)); //cap to max 10 lines
        float relativeH = linearNormalize(h, hBounds);
        m.put("%h", relativeH); //normalized line height

        //word features:
        String tok = elems.get(i).getPdfToken().token;

        getCaseMasks(tok).forEach(
          (String s) -> m.put(s, 1.0)); //case masks
        if (isStopWord(tok)) {
          m.put("%stop", 1.0); //stop word
          if (line != prevLine && (m.containsKey("%XXX") || m.containsKey("%Xxx")))
            m.put("%startCapStop", 1.0); //is a stop word that starts with a capital letter
        } else {
          if (m.containsKey("%xxx")) {
            m.put("%uncapns", 1.0); //is an uncapitalized stop word
          }
        }
        double adjLen = Math.min(tok.length(), 10.0) / 10.0;
        double adjLenSq = (adjLen - 0.5) * (adjLen - 0.5);
        m.put("%adjLen", adjLen); //adjusted word length
        m.put("%adjLenSq", adjLenSq); //adjusted word length squared (?)
        if (line <= 2)
          m.put("%first3lines", 1.0); //is it in the first three lines:
        if (lmFeats != null) { //how well does token match title/author gazeetters
          m.put("%tfreq", smoothFreq(tok, this.lmFeats.titleBow));
          m.put("%tffreq", smoothFreq(tok, this.lmFeats.titleFirstBow));
          m.put("%tlfreq", smoothFreq(tok, this.lmFeats.titleLastBow));
          m.put("%afreq", smoothFreq(Parser.fixupAuthors(tok), this.lmFeats.authorBow));
          m.put("%affreq", smoothFreq(Parser.fixupAuthors(tok), this.lmFeats.authorFirstBow));
          m.put("%alfreq", smoothFreq(Parser.fixupAuthors(tok), this.lmFeats.authorLastBow));
          m.put("%bfreq", smoothFreq(tok, this.lmFeats.backgroundBow));
          m.put("%bafreq", smoothFreq(Parser.fixupAuthors(tok), this.lmFeats.backgroundBow));
        }

        // add the token itself as a feature
        final String token = StringUtils.normalize(elems.get(i).getPdfToken().token);
		m.put("%t=" + token, 1.0);

        if(token.equals("and") || token.equals(","))
          m.put("%and", 1.0);

        // add trigram features
        final String trigramSourceToken = token + "$";
        for(int j = 0; j <= trigramSourceToken.length() - 3; ++j) {
          final String trigram = trigramSourceToken.substring(j, j + 3);
          final String feature = "%tri=" + trigram;
          m.updateValue(feature, 0.0, d -> d + 1);
        }

        // add word embeddings
        try {
          final Iterator<Double> vector = word2vecSearcher.getRawVector(tok).iterator();
          int j = 0;
          while(vector.hasNext()) {
            final double value = vector.next();
            m.put(wordEmbeddingFeatureNames[j], value);
            j += 1;
          }
        } catch (final Searcher.UnknownWordException e) {
          // do nothing
        }
      }
      out.add(m);
    }

    // print extensive debug information
    if(log.isDebugEnabled()) {
      // calculate a hash out of the tokens, so we can match feature values from different runs
      // based on the tokens
      final String tokens =
              out.stream().map(features ->
                      features.
                              keysView().
                              select(key -> key.startsWith("%t=")).
                              collect(featureName -> featureName.substring(3)).
                              makeString("-")
              ).collect(Collectors.joining(" "));
      final String tokensHashPrefix = String.format("%x", tokens.hashCode());

      log.debug("{} CRF Input for {}", tokensHashPrefix, tokens);
      PrintFeaturizedCRFInput.stringsFromFeaturizedSeq(out, tokensHashPrefix).stream().forEach(log::debug);
    }

    return out;
  }

  @Override
  public List<ObjectDoubleMap<String>> edgePredicates(List<PaperToken> elems) {
    val out = new ArrayList<ObjectDoubleMap<String>>();
    for (int i = 0; i < elems.size() - 1; i++) {
      val odhm = new ObjectDoubleHashMap<String>();
      odhm.put("B", 1.0);
      out.add(odhm);
    }
    return out; //I don't really understand these things.
  }

  private interface TokenPropertySelector {
    float getProp(PaperToken t);
  }
}
