content: package zemberek.lm;
  
  import com.google.common.collect.Lists;
  import java.io.BufferedInputStream;
  import java.io.BufferedOutputStream;
  import java.io.DataInputStream;
  import java.io.DataOutputStream;
  import java.io.File;
  import java.io.FileInputStream;
  import java.io.FileOutputStream;
  import java.io.IOException;
  import java.io.RandomAccessFile;
  import java.nio.file.Path;
  import java.text.Collator;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Locale;
  import java.util.Map;
  import java.util.TreeMap;
  import java.util.stream.Collectors;
  import zemberek.core.io.SimpleTextReader;
  import zemberek.core.logging.Log;
  
  public class LmVocabulary {
  
    public static final String DEFAULT_SENTENCE_BEGIN_MARKER = "<s>";
    public static final String DEFAULT_SENTENCE_END_MARKER = "</s>";
    public static final String DEFAULT_UNKNOWN_WORD = "<unk>";
    private String unknownWord;
    private String sentenceStart;
    private String sentenceEnd;
    private List<String> vocabulary;
    private Map<String, Integer> vocabularyIndexMap = new HashMap<>();
  
    private int unknownWordIndex = -1;
    private int sentenceStartIndex = -1;
    private int sentenceEndIndex = -1;
  
    /**
     * Generates a vocabulary with given String array.
     *
     * @param vocabulary word array.
     */
    public LmVocabulary(String... vocabulary) {
      Lists.newArrayList(vocabulary);
      generateMap(Lists.newArrayList(vocabulary));
    }
  
    /**
     * Generates a vocabulary with given String List.
     *
     * @param vocabulary word list.
     */
    public LmVocabulary(List<String> vocabulary) {
      generateMap(vocabulary);
    }
  
    /**
     * Generates a vocabulary from a binary RandomAccessFile. first integer read from the file pointer
     * defines the vocabulary size. Rest is read in UTF. Constructor does not close the
     * RandomAccessFile
     *
     * @param raf binary vocabulary RandomAccessFile
     */
    private LmVocabulary(RandomAccessFile raf) throws IOException {
      int vocabLength = raf.readInt();
      List<String> vocabulary = new ArrayList<>(vocabLength);
      for (int i = 0; i < vocabLength; i++) {
        vocabulary.add(raf.readUTF());
      }
      generateMap(vocabulary);
    }
  
    /**
     * Generates a vocabulary from a binary DataInputStream. first integer read from the file pointer
     * defines the vocabulary size. Rest is read in UTF. Constructor does not close the
     * DataInputStream
     *
     * @param dis input stream to read the vocabulary data.
     */
    private LmVocabulary(DataInputStream dis) throws IOException {
      loadVocabulary(dis);
    }
  
    /**
     * Generates a vocabulary from a binary vocabulary File. First integer in the file defines the
     * vocabulary size. Rest is read in UTF.
     */
    public static LmVocabulary loadFromBinary(File binaryVocabularyFile) throws IOException {
      try (DataInputStream dis = new DataInputStream(
          new BufferedInputStream(new FileInputStream(binaryVocabularyFile)))) {
        return new LmVocabulary(dis);
      }
    }
  
    /**
     * Generates a vocabulary from a binary vocabulary File. First integer in the file defines the
     * vocabulary size. Rest is read in UTF.
     */
    public static LmVocabulary loadFromBinary(Path binaryVocabularyFilePath) throws IOException {
      return loadFromBinary(binaryVocabularyFilePath.toFile());
    }
  
    /**
     * Generates a vocabulary from a binary RandomAccessFile. first integer read from the file pointer
     * defines the vocabulary size. Rest is read in UTF. The RandomAccessFile will not be closed by
     * this method.
     *
     * @param raf binary vocabulary RandomAccessFile
     */
    public static LmVocabulary loadFromRandomAccessFile(RandomAccessFile raf) throws IOException {
      return new LmVocabulary(raf);
    }
  
    /**
     * Generates a vocabulary from a binary DataInputStream. first integer read from the file pointer
     * defines the vocabulary size. Rest is read in UTF. This method does not close the
     * DataInputStream
     *
     * @param dis input stream to read the vocabulary data.
     */
    public static LmVocabulary loadFromDataInputStream(DataInputStream dis) throws IOException {
      return new LmVocabulary(dis);
    }
  
    /**
     * Generates a vocabulary from a UTF8-encoded text file.
     *
     * @param utfVocabularyFile input utf8 file to read vocabulary data.
     * @return LMVocabulary instance.
     */
    public static LmVocabulary loadFromUtf8File(File utfVocabularyFile) throws IOException {
      return new LmVocabulary(SimpleTextReader.trimmingUTF8Reader(utfVocabularyFile).asStringList());
    }
  
    /**
     * Returns a mutable Builder class which can be used for generating an LmVocabulary object.
     */
    public static Builder builder() {
      return new Builder();
    }
  
    /**
     * Generates a new LmVocabulary instance that contains words that exist in both `v1` and `v2`
     * There is no guarantee that new vocabulary indexes will match with v1 or v2.
     */
    public static LmVocabulary intersect(LmVocabulary v1, LmVocabulary v2) {
      HashSet<String> ls = new HashSet<>(v1.vocabulary);
      List<String> intersection = new ArrayList<>(Math.min(v1.size(), v2.size()));
      intersection.addAll(new HashSet<>(v2.vocabulary)
          .stream()
          .filter(ls::contains)
          .collect(Collectors.toList()));
      return new LmVocabulary(intersection);
    }
  
    private void loadVocabulary(DataInputStream dis) throws IOException {
      int vocabularyLength = dis.readInt();
      List<String> vocabulary = new ArrayList<>(vocabularyLength);
      for (int i = 0; i < vocabularyLength; i++) {
        vocabulary.add(dis.readUTF());
      }
      generateMap(vocabulary);
    }
  
    /**
     * Binary serialization of the vocabulary.
     *
     * @param file to serialize.
     */
    public void saveBinary(File file) throws IOException {
      try (DataOutputStream dos = new DataOutputStream(
          new BufferedOutputStream(new FileOutputStream(file)))) {
        saveBinary(dos);
      }
    }
  
    /**
     * Binary serialization of the vocabulary.
     *
     * @param dos output stream to serialize.
     */
    public void saveBinary(DataOutputStream dos) throws IOException {
      dos.writeInt(vocabulary.size());
      for (String s : vocabulary) {
        dos.writeUTF(s);
      }
    }
  
    private void generateMap(List<String> inputVocabulary) {
      // construct vocabulary index lookup.
      int indexCounter = 0;
      List<String> cleanVocab = new ArrayList<>();
      for (String word : inputVocabulary) {
        if (vocabularyIndexMap.containsKey(word)) {
          Log.warn("Language model vocabulary has duplicate item: " + word);
          continue;
        }
        if (word.equalsIgnoreCase(DEFAULT_UNKNOWN_WORD)) {
          if (unknownWordIndex != -1) {
            Log.warn(
                "Unknown word was already defined as %s but another matching token exist in the input vocabulary: %s",
                unknownWord, word);
          } else {
            unknownWord = word;
            unknownWordIndex = indexCounter;
          }
        } else if (word.equalsIgnoreCase(DEFAULT_SENTENCE_BEGIN_MARKER)) {
          if (sentenceStartIndex != -1) {
            Log.warn(
                "Sentence start index was already defined as %s but another matching token exist in the input vocabulary: %s",
                sentenceStart, word);
          } else {
            sentenceStart = word;
            sentenceStartIndex = indexCounter;
          }
        } else if (word.equalsIgnoreCase(DEFAULT_SENTENCE_END_MARKER)) {
          if (sentenceEndIndex != -1) {
            Log.warn(
                "Sentence end index was already defined as %s but another matching token exist in the input vocabulary: %s",
                sentenceEnd, word);
          } else {
            sentenceEnd = word;
            sentenceEndIndex = indexCounter;
          }
        }
        vocabularyIndexMap.put(word, indexCounter);
        cleanVocab.add(word);
        indexCounter++;
      }
      if (unknownWordIndex == -1) {
        unknownWord = DEFAULT_UNKNOWN_WORD;
        cleanVocab.add(unknownWord);
        vocabularyIndexMap.put(unknownWord, indexCounter++);
        Log.debug("Necessary special token " + unknownWord
            + " was not found in the vocabulary, it is added explicitly");
      }
      unknownWordIndex = vocabularyIndexMap.get(unknownWord);
      if (sentenceStartIndex == -1) {
        sentenceStart = DEFAULT_SENTENCE_BEGIN_MARKER;
        cleanVocab.add(sentenceStart);
        vocabularyIndexMap.put(sentenceStart, indexCounter++);
        Log.debug("Vocabulary does not contain sentence start token, it is added explicitly.");
      }
      sentenceStartIndex = vocabularyIndexMap.get(sentenceStart);
      if (sentenceEndIndex == -1) {
        sentenceEnd = DEFAULT_SENTENCE_END_MARKER;
        cleanVocab.add(sentenceEnd);
        vocabularyIndexMap.put(sentenceEnd, indexCounter);
        Log.debug("Vocabulary does not contain sentence end token, it is added explicitly.");
      }
      sentenceEndIndex = vocabularyIndexMap.get(sentenceEnd);
      vocabulary = Collections.unmodifiableList(cleanVocab);
    }
  
    public int size() {
      // Because we may have duplicate items in word list but not in map, we return the size of the list.
      return vocabulary.size();
    }
  
    public boolean containsUnknown(int... gramIds) {
      for (int gramId : gramIds) {
        if (gramId < 0 || gramId >= vocabulary.size() || gramId == unknownWordIndex) {
          return true;
        }
      }
      return false;
    }
  
    /**
     * returns true if any word in vocabulary starts with `_` or `-`
     */
    public boolean containsSuffix() {
      for (String s : vocabulary) {
        if (s.startsWith("_") || s.startsWith("-")) {
          return true;
        }
      }
      return false;
    }
  
    /**
     * returns true if any word in vocabulary ends with `_` or `-`
     */
    public boolean containsPrefix() {
      for (String s : vocabulary) {
        if (s.endsWith("_") || s.endsWith("-")) {
          return true;
        }
      }
      return false;
    }
  
    /**
     * @param index Word for the index. if index is out of bounds, <UNK> is returned with warning.
     * Note that Vocabulary may contain <UNK> token as well.
     */
    public String getWord(int index) {
      if (index < 0 || index >= vocabulary.size()) {
        Log.warn("Out of bounds word index is used:" + index);
        return unknownWord;
      }
      return vocabulary.get(index);
    }
  
    public int indexOf(String word) {
      Integer k = vocabularyIndexMap.get(word);
      return k == null ? unknownWordIndex : k;
    }
  
    public int getSentenceStartIndex() {
      return sentenceStartIndex;
    }
  
    public int getSentenceEndIndex() {
      return sentenceEndIndex;
    }
  
    public int getUnknownWordIndex() {
      return unknownWordIndex;
    }
  
    public String getUnknownWord() {
      return unknownWord;
    }
  
    public String getSentenceStart() {
      return sentenceStart;
    }
  
    public String getSentenceEnd() {
      return sentenceEnd;
    }
  
    /**
     * @return indexes of words when the words are alphabetically sorted according to the input
     * locale.
     */
    public Iterable<Integer> alphabeticallySortedWordsIds(Locale locale) {
      return new TreeMap<String, Integer>(Collator.getInstance(locale)).values();
    }
  
    /**
     * @return indexes of words when the words are alphabetically sorted according to the default
     * locale.
     */
    public Iterable<Integer> alphabeticallySortedWordsIds() {
      return new TreeMap<>(vocabularyIndexMap).values();
    }
  
    public Iterable<String> words() {
      return this.vocabulary;
    }
  
    public Iterable<String> wordsSorted() {
      List<String> sorted = new ArrayList<>(vocabulary);
      Collections.sort(sorted);
      return sorted;
    }
  
    public Iterable<String> wordsSorted(Locale locale) {
      List<String> sorted = new ArrayList<>(vocabulary);
      sorted.sort(Collator.getInstance(locale));
      return sorted;
    }
  
    /**
     * @param indexes word indexes
     * @return the Word representation of indexes. Such as for 2,3,4 it returns "foo bar zipf" For
     * unknown indexes it uses <unk>.
     */
    public String getWordsString(int... indexes) {
      StringBuilder sb = new StringBuilder();
      for (int i = 0; i < indexes.length; i++) {
        int index = indexes[i];
        if (contains(index)) {
          sb.append(vocabulary.get(index));
        } else {
          Log.warn("Out of bounds word index is used:" + index);
          sb.append(unknownWord);
        }
        if (i < indexes.length - 1) {
          sb.append(" ");
        }
      }
      return sb.toString();
    }
  
    /**
     * @param indexes index array
     * @return true if all indexes are within the Vocabulary boundaries.
     */
    public boolean containsAll(int... indexes) {
      for (int index : indexes) {
        if (!contains(index)) {
          return false;
        }
      }
      return true;
    }
  
    /**
     * @param index word index
     * @return true if index is within the Vocabulary boundaries.
     */
    public boolean contains(int index) {
      return index >= 0 && index < vocabulary.size();
    }
  
    /**
     * @param words Words
     * @return true if vocabulary contains all words.
     */
    public boolean containsAll(String... words) {
      for (String word : words) {
        if (!contains(word)) {
          return false;
        }
      }
      return true;
    }
  
    /**
     * @param word Word
     * @return if vocabulary contains the word. For special tokens, it always return true.
     */
    public boolean contains(String word) {
      return vocabularyIndexMap.containsKey(word);
    }
  
    /**
     * @param g0 W0 Index
     * @param g1 W1 Index
     * @param g2 W2 Index
     * @return The encoded long representation of a trigram. Structure: [1 bit EMPTY][21 bit
     * W2-IND][21 bit W1-IND][21 bit W0-IND]
     */
    public long encodeTrigram(int g0, int g1, int g2) {
      long encoded = g2;
      encoded = (encoded << 21) | g1;
      return (encoded << 21) | g0;
    }
  
    /**
     * @param triGram trigram Indexes.
     * @return the encoded long representation of a trigram. Structure: [1 bit EMPTY][21 bit
     * W2-IND][21 bit W1-IND][21 bit W0-IND]
     */
    public long encodeTrigram(int... triGram) {
      if (triGram.length > 3) {
        throw new IllegalArgumentException(
            "Cannot generate long from order " + triGram.length + " grams ");
      }
      long encoded = triGram[2];
      encoded = (encoded << 21) | triGram[1];
      return (encoded << 21) | triGram[0];
    }
  
    /**
     * @param words word array
     * @return the vocabulary index array for a word array. if a word is unknown, index of <UNK> is
     * used is returned as its vocabulary index. This value can be -1.
     */
    public int[] toIndexes(String... words) {
      int[] indexes = new int[words.length];
      int i = 0;
      for (String word : words) {
        if (!vocabularyIndexMap.containsKey(word)) {
          indexes[i] = unknownWordIndex;
        } else {
          indexes[i] = vocabularyIndexMap.get(word);
        }
        i++;
      }
      return indexes;
    }
  
    /**
     * @param history and current word.
     * @return the vocabulary index array for a word array. if a word is unknown, index of <UNK> is
     * used is returned as its vocabulary index. This value can be -1.
     */
    public int[] toIndexes(String[] history, String word) {
      int[] indexes = new int[history.length + 1];
      for (int j = 0; j <= history.length; j++) {
        String s = j < history.length ? history[j] : word;
        if (!vocabularyIndexMap.containsKey(s)) {
          indexes[j] = unknownWordIndex;
        } else {
          indexes[j] = vocabularyIndexMap.get(s);
        }
      }
      return indexes;
    }
  
    /**
     * word index history and current word.
     *
     * @return the vocabulary index array for a word array. if a word is unknown, index of <UNK> is
     * used is returned as its vocabulary index. This value can be -1.
     */
    public int[] toIndexes(int[] history, String word) {
      int[] indexes = new int[history.length + 1];
      for (int j = 0; j <= history.length; j++) {
        int index = j < history.length ? history[j] : indexOf(word);
        if (!contains(index)) {
          indexes[j] = unknownWordIndex;
        } else {
          indexes[j] = index;
        }
      }
      return indexes;
    }
  
    /**
     * @param indexes word indexes.
     * @return Words representations of the indexes. If an index is out of bounds, <UNK>
     * representation is used. Note that Vocabulary may contain an <UNK> word in it.
     */
    public String[] toWords(int... indexes) {
      String[] words = new String[indexes.length];
      int k = 0;
      for (int index : indexes) {
        if (contains(index)) {
          words[k++] = vocabulary.get(index);
        } else {
          Log.warn("Out of bounds word index is used:" + index);
          words[k++] = unknownWord;
        }
      }
      return words;
    }
  
    /**
     * This class acts like a mutable vocabulary. It can be used for dynamically generating an
     * LmVocabulary object.
     */
    public static class Builder {
  
      private Map<String, Integer> map = new HashMap<>();
      private List<String> tokens = Lists.newArrayList();
  
      public int add(String word) {
        int index = indexOf(word);
        if (index != -1) {
          return index;
        } else {
          index = tokens.size();
          map.put(word, index);
          tokens.add(word);
          return index;
        }
      }
  
      int size() {
        return tokens.size();
      }
  
      public Builder addAll(String... words) {
        for (String word : words) {
          add(word);
        }
        return this;
      }
  
      public Builder addAll(Iterable<String> words) {
        for (String word : words) {
          add(word);
        }
        return this;
      }
  
      /**
       * @return index of input. -1 if input does not exist.
       */
      public int indexOf(String key) {
        return map.getOrDefault(key, -1);
      }
  
      /**
       * @return an unmodifiable copy of the words so far added.
       */
      public List<String> words() {
        return Collections.unmodifiableList(tokens);
      }
  
      /**
       * @return indexes of words when the words are alphabetically sorted according to the input
       * locale.
       */
      public Iterable<Integer> alphabeticallySortedWordsIds(Locale locale) {
        TreeMap<String, Integer> treeMap = new TreeMap<>(Collator.getInstance(locale));
        treeMap.putAll(map);
        return treeMap.values();
      }
  
      /**
       * @return indexes of words when the words are alphabetically sorted according to EN locale.
       */
      public Iterable<Integer> alphabeticallySortedWordsIds() {
        return new TreeMap<>(map).values();
      }
  
      /**
       * @return Generated unmodifiable LmVocabulary
       */
      public LmVocabulary generate() {
        return new LmVocabulary(tokens);
      }
  
    }
  }