<commit_before>/**
   Copyright (c) 2007, University of Pittsburgh

   All rights reserved.

   Redistribution and use in source and binary forms, with or without
   modification, are permitted provided that the following conditions are
   met:

 * Redistributions of source code must retain the above copyright
   notice, this list of conditions and the following disclaimer.

 * Redistributions in binary form must reproduce the above
   copyright notice, this list of conditions and the following
   disclaimer in the documentation and/or other materials provided
   with the distribution.

 * Neither the name of the University of Pittsburgh nor the names
   of its contributors may be used to endorse or promote products
   derived from this software without specific prior written
   permission.

   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
   "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
   A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
   CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
   PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
   LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
   NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 **/

package pitt.search.semanticvectors;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.nio.file.FileSystems;
import java.util.ArrayList;
import java.util.Hashtable;
import java.util.List;
import java.util.TreeSet;
import java.util.logging.Logger;

import org.apache.lucene.document.Document;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.BaseCompositeReader;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.DocsEnum;
import org.apache.lucene.index.FieldInfos;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.MultiFields;
import org.apache.lucene.index.SlowCompositeReaderWrapper;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.Terms;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.Version;

import pitt.search.semanticvectors.utils.StringUtils;
import pitt.search.semanticvectors.utils.VerbatimLogger;

/**
 * Class to support reading extra information from Lucene indexes,
 * including term frequency, doc frequency.
 */
public class LuceneUtils {
  public static final Version LUCENE_VERSION = Version.LUCENE_5_0_0;

  private static final Logger logger = Logger.getLogger(DocVectors.class.getCanonicalName());
  private FlagConfig flagConfig;
  private BaseCompositeReader<LeafReader> compositeReader;
  private LeafReader leafReader;
  private Hashtable<Term, Float> termEntropy = new Hashtable<Term, Float>();
  private Hashtable<Term, Float> termIDF = new Hashtable<>();
  private TreeSet<String> stopwords = null;
  private TreeSet<String> startwords = null;

  /**
   * Determines which term-weighting strategy to use in indexing, 
   * and in search if {@link FlagConfig#usetermweightsinsearch()} is set.
   * 
   * <p>Names may be passed as command-line arguments, so underscores are avoided.
   */
  public enum TermWeight {
    /** No term weighting: all terms have weight 1. */
    NONE,
    /** Use inverse document frequency: see {@link LuceneUtils#getIDF}. */
    IDF,
    /** Use log entropy: see {@link LuceneUtils#getEntropy}. */
    LOGENTROPY,
    /** Use square root of term frequency. */
    SQRT,
  }

  /**
   * @param flagConfig Contains all information necessary for configuring LuceneUtils.
   *        {@link FlagConfig#luceneindexpath()} must be non-empty. 
   */
  public LuceneUtils(FlagConfig flagConfig) throws IOException {
    if (flagConfig.luceneindexpath().isEmpty()) {
      throw new IllegalArgumentException(
          "-luceneindexpath is a required argument for initializing LuceneUtils instance.");
    }

    this.compositeReader = DirectoryReader.open(
        FSDirectory.open(FileSystems.getDefault().getPath(flagConfig.luceneindexpath())));
    this.leafReader =  SlowCompositeReaderWrapper.wrap(compositeReader);
    MultiFields.getFields(compositeReader);
    this.flagConfig = flagConfig;
    if (!flagConfig.stoplistfile().isEmpty())
      loadStopWords(flagConfig.stoplistfile());

    if (!flagConfig.startlistfile().isEmpty())
        loadStartWords(flagConfig.startlistfile());

    VerbatimLogger.info("Initialized LuceneUtils from Lucene index in directory: " + flagConfig.luceneindexpath() + "\n");
  }

  /**
   * Loads the stopword file into the {@link #stopwords} data structure.
   * @param stoppath Path to stopword file.
   * @throws IOException If stopword file cannot be read.
   */
  public void loadStopWords(String stoppath) throws IOException  {
    logger.info("Using stopword file: " + stoppath);
    stopwords = new TreeSet<String>();
    try {
      BufferedReader readIn = new BufferedReader(new FileReader(stoppath));
      String in = readIn.readLine();
      while (in != null) {
        stopwords.add(in);
        in = readIn.readLine();
      }
      readIn.close();
    }
    catch (IOException e) {
      throw new IOException("Couldn't open file " + stoppath);
    }
  }

  /**
   * Loads the startword file into the {@link #startwords} data structure.
   * @param startpath Path to startword file
   * @throws IOException If startword file cannot be read.
   */
  public void loadStartWords(String startpath) throws IOException  {
    startwords = new TreeSet<String>();
    try {
      BufferedReader readIn = new BufferedReader(new FileReader(startpath));
      String in = readIn.readLine();
      while (in != null) {
        startwords.add(in);
        in = readIn.readLine();
      }
      VerbatimLogger.info(String.format(
          "Loading startword file: '%s'. Only these %d words will be indexed.\n",
          startpath, startwords.size()));
      readIn.close();
    }
    catch (IOException e) {
      throw new IOException("Couldn't open file "+startpath);
    }
  }

  /**
   * Returns true if term is in stoplist, false otherwise.
   */
  public boolean stoplistContains(String x) {
    if (stopwords == null) return false;
    return stopwords.contains(x);
  }
  
  /**
   * Returns false if term is not in startlist, true otherwise (including if no startlist exists).
   */
  public boolean startlistContains(String x) {
    if (startwords == null) return true;
    return startwords.contains(x);
  }
  
  public Document getDoc(int docID) throws IOException {
    return this.leafReader.document(docID);
  }

  public String getExternalDocId(int docID) throws IOException {
    String externalDocId;
    try {
      externalDocId = this.getDoc(docID).getField(flagConfig.docidfield()).stringValue();
    } catch (IOException e) {
      logger.severe(String.format(
          "Failed to get external doc ID from doc no. %d in Lucene index." +
              "\nThis is almost certain to lead to problems." +
              "\nCheck that -docidfield was set correctly and exists in the Lucene index",
          docID));
      throw e;
    }

    return externalDocId;
  }

  /**
   * Gets the terms for a given field. Throws {@link java.lang.NullPointerException} if this is null.<commit_msg>Added option to avoid looking up external document ID if so desired (to save time, or if no suitable field exists)<commit_after>/**
   Copyright (c) 2007, University of Pittsburgh

   All rights reserved.

   Redistribution and use in source and binary forms, with or without
   modification, are permitted provided that the following conditions are
   met:

 * Redistributions of source code must retain the above copyright
   notice, this list of conditions and the following disclaimer.

 * Redistributions in binary form must reproduce the above
   copyright notice, this list of conditions and the following
   disclaimer in the documentation and/or other materials provided
   with the distribution.

 * Neither the name of the University of Pittsburgh nor the names
   of its contributors may be used to endorse or promote products
   derived from this software without specific prior written
   permission.

   THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
   "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
   LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
   A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
   CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
   PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
   PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
   LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
   NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
   SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 **/

package pitt.search.semanticvectors;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.nio.file.FileSystems;
import java.util.ArrayList;
import java.util.Hashtable;
import java.util.List;
import java.util.TreeSet;
import java.util.logging.Logger;

import org.apache.lucene.document.Document;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.BaseCompositeReader;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.DocsEnum;
import org.apache.lucene.index.FieldInfos;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.MultiFields;
import org.apache.lucene.index.SlowCompositeReaderWrapper;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.Terms;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.Version;

import pitt.search.semanticvectors.utils.StringUtils;
import pitt.search.semanticvectors.utils.VerbatimLogger;

/**
 * Class to support reading extra information from Lucene indexes,
 * including term frequency, doc frequency.
 */
public class LuceneUtils {
  public static final Version LUCENE_VERSION = Version.LUCENE_5_0_0;

  private static final Logger logger = Logger.getLogger(DocVectors.class.getCanonicalName());
  private FlagConfig flagConfig;
  private BaseCompositeReader<LeafReader> compositeReader;
  private LeafReader leafReader;
  private Hashtable<Term, Float> termEntropy = new Hashtable<Term, Float>();
  private Hashtable<Term, Float> termIDF = new Hashtable<>();
  private TreeSet<String> stopwords = null;
  private TreeSet<String> startwords = null;

  /**
   * Determines which term-weighting strategy to use in indexing, 
   * and in search if {@link FlagConfig#usetermweightsinsearch()} is set.
   * 
   * <p>Names may be passed as command-line arguments, so underscores are avoided.
   */
  public enum TermWeight {
    /** No term weighting: all terms have weight 1. */
    NONE,
    /** Use inverse document frequency: see {@link LuceneUtils#getIDF}. */
    IDF,
    /** Use log entropy: see {@link LuceneUtils#getEntropy}. */
    LOGENTROPY,
    /** Use square root of term frequency. */
    SQRT,
  }

  /**
   * @param flagConfig Contains all information necessary for configuring LuceneUtils.
   *        {@link FlagConfig#luceneindexpath()} must be non-empty. 
   */
  public LuceneUtils(FlagConfig flagConfig) throws IOException {
    if (flagConfig.luceneindexpath().isEmpty()) {
      throw new IllegalArgumentException(
          "-luceneindexpath is a required argument for initializing LuceneUtils instance.");
    }

    this.compositeReader = DirectoryReader.open(
        FSDirectory.open(FileSystems.getDefault().getPath(flagConfig.luceneindexpath())));
    this.leafReader =  SlowCompositeReaderWrapper.wrap(compositeReader);
    MultiFields.getFields(compositeReader);
    this.flagConfig = flagConfig;
    if (!flagConfig.stoplistfile().isEmpty())
      loadStopWords(flagConfig.stoplistfile());

    if (!flagConfig.startlistfile().isEmpty())
        loadStartWords(flagConfig.startlistfile());

    VerbatimLogger.info("Initialized LuceneUtils from Lucene index in directory: " + flagConfig.luceneindexpath() + "\n");
  }

  /**
   * Loads the stopword file into the {@link #stopwords} data structure.
   * @param stoppath Path to stopword file.
   * @throws IOException If stopword file cannot be read.
   */
  public void loadStopWords(String stoppath) throws IOException  {
    logger.info("Using stopword file: " + stoppath);
    stopwords = new TreeSet<String>();
    try {
      BufferedReader readIn = new BufferedReader(new FileReader(stoppath));
      String in = readIn.readLine();
      while (in != null) {
        stopwords.add(in);
        in = readIn.readLine();
      }
      readIn.close();
    }
    catch (IOException e) {
      throw new IOException("Couldn't open file " + stoppath);
    }
  }

  /**
   * Loads the startword file into the {@link #startwords} data structure.
   * @param startpath Path to startword file
   * @throws IOException If startword file cannot be read.
   */
  public void loadStartWords(String startpath) throws IOException  {
    startwords = new TreeSet<String>();
    try {
      BufferedReader readIn = new BufferedReader(new FileReader(startpath));
      String in = readIn.readLine();
      while (in != null) {
        startwords.add(in);
        in = readIn.readLine();
      }
      VerbatimLogger.info(String.format(
          "Loading startword file: '%s'. Only these %d words will be indexed.\n",
          startpath, startwords.size()));
      readIn.close();
    }
    catch (IOException e) {
      throw new IOException("Couldn't open file "+startpath);
    }
  }

  /**
   * Returns true if term is in stoplist, false otherwise.
   */
  public boolean stoplistContains(String x) {
    if (stopwords == null) return false;
    return stopwords.contains(x);
  }
  
  /**
   * Returns false if term is not in startlist, true otherwise (including if no startlist exists).
   */
  public boolean startlistContains(String x) {
    if (startwords == null) return true;
    return startwords.contains(x);
  }
  
  public Document getDoc(int docID) throws IOException {
    return this.leafReader.document(docID);
  }

  public String getExternalDocId(int docID) throws IOException {
    
	  //to save time, avoid using external ID if so desired
	  if (flagConfig.docidfield().equals("luceneID")) return docID +"";
	  
	String externalDocId;
    try {
      externalDocId = this.getDoc(docID).getField(flagConfig.docidfield()).stringValue();
    } catch (IOException e) {
      logger.severe(String.format(
          "Failed to get external doc ID from doc no. %d in Lucene index." +
              "\nThis is almost certain to lead to problems." +
              "\nCheck that -docidfield was set correctly and exists in the Lucene index",
          docID));
      throw e;
    }

    return externalDocId;
  }

  /**
   * Gets the terms for a given field. Throws {@link java.lang.NullPointerException} if this is null.
{'fraction_non_alphanumeric': 0.08894090111176127, 'fraction_numerical': 0.0010402444574475, 'mean_word_length': 3.616446578631453, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5042141', 'n_tokens_mistral': 4414, 'n_tokens_neox': 4210, 'n_words': 1447}