package com.pera.suggestion;

import java.io.IOException;

import java.io.Reader;

import java.util.ArrayList;

import java.util.HashMap;

import java.util.Iterator;

import java.util.List;

import java.util.Map;

import org.apache.lucene.analysis.Analyzer;

import org.apache.lucene.analysis.LowerCaseFilter;

import org.apache.lucene.analysis.StopFilter;

import org.apache.lucene.analysis.TokenStream;

import org.apache.lucene.analysis.standard.StandardFilter;

import org.apache.lucene.analysis.standard.StandardTokenizer;

import org.apache.lucene.document.Document;

import org.apache.lucene.document.Field;

import org.apache.lucene.index.CorruptIndexException;

import org.apache.lucene.index.IndexReader;

import org.apache.lucene.index.IndexWriter;

import org.apache.lucene.index.Term;

import org.apache.lucene.search.IndexSearcher;

import org.apache.lucene.search.Query;

import org.apache.lucene.search.ScoreDoc;

import org.apache.lucene.search.Sort;

import org.apache.lucene.search.TermQuery;

import org.apache.lucene.search.TopDocs;

import org.apache.lucene.store.Directory;

import org.apache.lucene.store.FSDirectory;

public class Sugesstion {

private static final String GRAMMED_WORDS_FIELD = "words";

private static final String SOURCE_WORD_FIELD = "sourceWord";

private static final String COUNT_FIELD = "count";

private static final String[] ENGLISH_STOP_WORDS = {

     "a", "an", "and", "are", "as", "at", "be", "but", "by",

     "for", "i", "if", "in", "into", "is",

     "no", "not", "of", "on", "or", "s", "such",

     "t", "that", "the", "their", "then", "there", "these",

     "they", "this", "to", "was", "will", "with"

     };

private final Directory autoCompleteDirectory;

private IndexReader autoCompleteReader;

private IndexSearcher autoCompleteSearcher;

public Sugesstion(String autoCompleteDir) throws IOException {

      this.autoCompleteDirectory = FSDirectory.getDirectory(autoCompleteDir,

        null);

reOpenReader();

     }

public List<String> suggestTermsFor(String term) throws IOException {

      // get the top 5 terms for query

      Query query = new TermQuery(new Term(GRAMMED_WORDS_FIELD, term));

      Sort sort = new Sort(COUNT_FIELD, true);

TopDocs docs = autoCompleteSearcher.search(query, null, 5, sort);

      List<String> suggestions = new ArrayList<String>();

      for (ScoreDoc doc : docs.scoreDocs) {

       suggestions.add(autoCompleteReader.document(doc.doc).get(

         SOURCE_WORD_FIELD));

      }

return suggestions;

     }

@SuppressWarnings("unchecked")

     public void reIndex(Directory sourceDirectory, String fieldToAutocomplete)

       throws CorruptIndexException, IOException {

      // build a dictionary (from the spell package)

      IndexReader sourceReader = IndexReader.open(sourceDirectory);

LuceneDictionary dict = new LuceneDictionary(sourceReader,

        fieldToAutocomplete);

// code from

      // org.apache.lucene.search.spell.SpellChecker.indexDictionary(

      // Dictionary)

      IndexReader.unlock(autoCompleteDirectory);

// use a custom analyzer so we can do EdgeNGramFiltering

      IndexWriter writer = new IndexWriter(autoCompleteDirectory,

      new Analyzer() {

       public TokenStream tokenStream(String fieldName,

         Reader reader) {

        TokenStream result = new StandardTokenizer(reader);

result = new StandardFilter(result);

        result = new LowerCaseFilter(result);

        result = new ISOLatin1AccentFilter(result);

        result = new StopFilter(result,

         ENGLISH_STOP_WORDS);

        result = new EdgeNGramTokenFilter(

         result, Side.FRONT,1, 20);

return result;

       }

      }, true);

writer.setMergeFactor(300);

      writer.setMaxBufferedDocs(150);

// go through every word, storing the original word (incl. n-grams)

      // and the number of times it occurs

      Map<String, Integer> wordsMap = new HashMap<String, Integer>();

Iterator<String> iter = (Iterator<String>) dict.getWordsIterator();

      while (iter.hasNext()) {

       String word = iter.next();

int len = word.length();

       if (len < 3) {

        continue; // too short we bail but "too long" is fine...

       }

if (wordsMap.containsKey(word)) {

        throw new IllegalStateException(

          "This should never happen in Lucene 2.3.2");

        // wordsMap.put(word, wordsMap.get(word) + 1);

       } else {

        // use the number of documents this word appears in

        wordsMap.put(word, sourceReader.docFreq(new Term(

          fieldToAutocomplete, word)));

       }

      }

for (String word : wordsMap.keySet()) {

       // ok index the word

       Document doc = new Document();

       doc.add(new Field(SOURCE_WORD_FIELD, word, Field.Store.YES,

         Field.Index.UN_TOKENIZED)); // orig term

       doc.add(new Field(GRAMMED_WORDS_FIELD, word, Field.Store.YES,

         Field.Index.TOKENIZED)); // grammed

       doc.add(new Field(COUNT_FIELD,

         Integer.toString(wordsMap.get(word)), Field.Store.NO,

         Field.Index.UN_TOKENIZED)); // count

writer.addDocument(doc);

      }

sourceReader.close();

// close writer

      writer.optimize();

      writer.close();

// re-open our reader

      reOpenReader();

     }

private void reOpenReader() throws CorruptIndexException, IOException {

      if (autoCompleteReader == null) {

       autoCompleteReader = IndexReader.open(autoCompleteDirectory);

      } else {

       autoCompleteReader.reopen();

      }

autoCompleteSearcher = new IndexSearcher(autoCompleteReader);

     }

public static void main(String[] args) throws Exception {

      Sugesstion autocomplete = new Sugesstion("/index/autocomplete");

// run this to re-index from the current index, shouldn't need to do

      // this very often

      // autocomplete.reIndex(FSDirectory.getDirectory("/index/live", null),

      // "content");

String term = "steve";

System.out.println(autocomplete.suggestTermsFor(term));

      // prints [steve, steven, stevens, stevenson, stevenage]

     }

}