// Generalized by: ElementInstanceMatrix

// A represents the id of a document to be indexed (e.g. a snippet ID)
transient sclass WordIndexWithBitSets<A> {
  S regexp = "\\w+";
  new UniqueList<A> documents;
  Map<S, Entry> index = ciMap();

  // each word can be represented as either an int array, a bit set or both
  // for now, the int array is always there
  sclass Entry {
    new Set<Int> intSet; // while building the index
    int[] intArray;
    BitSet bitSet;
    
    BitSet bitSet() {
      if (bitSet == null)
        bitSet = intArrayToBitSet(intArray);
      ret bitSet;
    }
    
    int count() { ret intArray.length; }
  }
  
  *() {}
  *(Map<A, S> map) { fOr (A a, S text : map) add(a, text); }
  
  void add(A a, S text) {
    int idx = documents.addOrGetIndex(a);
    Set<S> words = extractWords(text);
    for (S word : words) addWord(idx, word);
  }
  
  void addWord(int idx, S word) {
    Entry e = index.get(word);
    if (e == null) index.put(word, e = new Entry);
    e.intSet.add(idx);
  }
  
  // Call this exactly once before doing queries on the index
  void doneAdding() {
    for (Entry e : values(index)) {
      e.intArray = sortIntArrayInPlace(toIntArray(e.intSet));
      e.intSet = null;
    }
  }
  
  Set<S> extractWords(S text) {
    ret asCISet(extractWords_list(text));
  }
  
  LS extractWords_list(S text) {
    ret regexpExtractAll(regexp, text);
  }
  
  L<IntRange> wordRanges(S text) {
    ret regexpFindRanges(regexp, text);
  }
  
  L<A> documentsContainingWord(S word) {
    Entry e = index.get(word);
    ret e == null ? null : documentListFromIntArray(e.intArray);
  }
  
  L<A> documentListFromIntArray(int[] array) {
    ret listFromFunction(array.length, i -> documents.get(array[i]));
  }
  
  // may return null (=empty list)
  L<A> documentsContainingAllWords(Cl<S> words) {
    int n = l(words;
    if (n == 0) ret documents;
    if (n == 1) ret documentsContainingWord(first(words));

    // get entries for words, exit when a word is unknown
    Entry[] entries = new[n];
    Iterator<S> itWords = iterator(words);
    for i to n: {
      Entry e = index.get(itWords.next());
      if (e == null) null;
      entries[i] = e;
    }
    
    // go through words again, AND-combine all bit sets
    BitSet bs = cloneBitSet(entries[0].bitSet());
    for (int i = 1; i < n; i++)
      bs.and(entries[i].bitSet());

    ret documentListFromIntArray(bitSetToIntArray(bs)); 
  }
  
  L<A> documentsContainingAllWords(S text, O... _) {
    ret documentsContainingAllWords(extractWords(text));
  }
  
  NavigableSet<S> words() { ret (NavigableSet) keys(index); }
  
  int numWords() { ret l(index); }
  
  // This method only works when A = S
  void add(S s) { add((A) s, s); }
}