content: package org.apache.lucene.search.highlight;
  
  /**
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  import java.io.IOException;
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.Iterator;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  
  import org.apache.lucene.analysis.CachingTokenFilter;
  import org.apache.lucene.analysis.TokenStream;
  import org.apache.lucene.index.FilterIndexReader;
  import org.apache.lucene.index.IndexReader;
  import org.apache.lucene.index.Term;
  import org.apache.lucene.index.TermEnum;
  import org.apache.lucene.index.memory.MemoryIndex;
  import org.apache.lucene.search.*;
  import org.apache.lucene.search.spans.FieldMaskingSpanQuery;
  import org.apache.lucene.search.spans.SpanFirstQuery;
  import org.apache.lucene.search.spans.SpanNearQuery;
  import org.apache.lucene.search.spans.SpanNotQuery;
  import org.apache.lucene.search.spans.SpanOrQuery;
  import org.apache.lucene.search.spans.SpanQuery;
  import org.apache.lucene.search.spans.SpanTermQuery;
  import org.apache.lucene.search.spans.Spans;
  import org.apache.lucene.util.StringHelper;
  
  /**
   * Class used to extract {@link WeightedSpanTerm}s from a {@link Query} based on whether 
   * {@link Term}s from the {@link Query} are contained in a supplied {@link TokenStream}.
   */
  public class WeightedSpanTermExtractor {
  
    private String fieldName;
    private TokenStream tokenStream;
    private Map<String,IndexReader> readers = new HashMap<String,IndexReader>(10); 
    private String defaultField;
    private boolean expandMultiTermQuery;
    private boolean cachedTokenStream;
    private boolean wrapToCaching = true;
    private int maxDocCharsToAnalyze;
  
    public WeightedSpanTermExtractor() {
    }
  
    public WeightedSpanTermExtractor(String defaultField) {
      if (defaultField != null) {
        this.defaultField = StringHelper.intern(defaultField);
      }
    }
  
    private void closeReaders() {
      Collection<IndexReader> readerSet = readers.values();
  
      for (final IndexReader reader : readerSet) {
        try {
          reader.close();
        } catch (IOException e) {
          // alert?
        }
      }
    }
  
    /**
     * Fills a <code>Map</code> with <@link WeightedSpanTerm>s using the terms from the supplied <code>Query</code>.
     * 
     * @param query
     *          Query to extract Terms from
     * @param terms
     *          Map to place created WeightedSpanTerms in
     * @throws IOException
     */
    private void extract(Query query, Map<String,WeightedSpanTerm> terms) throws IOException {
      if (query instanceof BooleanQuery) {
        BooleanClause[] queryClauses = ((BooleanQuery) query).getClauses();
  
        for (int i = 0; i < queryClauses.length; i++) {
          if (!queryClauses[i].isProhibited()) {
            extract(queryClauses[i].getQuery(), terms);
          }
        }
      } else if (query instanceof PhraseQuery) {
        PhraseQuery phraseQuery = ((PhraseQuery) query);
        Term[] phraseQueryTerms = phraseQuery.getTerms();
        SpanQuery[] clauses = new SpanQuery[phraseQueryTerms.length];
        for (int i = 0; i < phraseQueryTerms.length; i++) {
          clauses[i] = new SpanTermQuery(phraseQueryTerms[i]);
        }
        int slop = phraseQuery.getSlop();
        int[] positions = phraseQuery.getPositions();
        // add largest position increment to slop
        if (positions.length > 0) {
          int lastPos = positions[0];
          int largestInc = 0;
          int sz = positions.length;
          for (int i = 1; i < sz; i++) {
            int pos = positions[i];
            int inc = pos - lastPos;
            if (inc > largestInc) {
              largestInc = inc;
            }
            lastPos = pos;
          }
          if(largestInc > 1) {
            slop += largestInc;
          }
        }
  
        boolean inorder = false;
  
        if (slop == 0) {
          inorder = true;
        }
  
        SpanNearQuery sp = new SpanNearQuery(clauses, slop, inorder);
        sp.setBoost(query.getBoost());
        extractWeightedSpanTerms(terms, sp);
      } else if (query instanceof TermQuery) {
        extractWeightedTerms(terms, query);
      } else if (query instanceof SpanQuery) {
        extractWeightedSpanTerms(terms, (SpanQuery) query);
      } else if (query instanceof FilteredQuery) {
        extract(((FilteredQuery) query).getQuery(), terms);
      } else if (query instanceof DisjunctionMaxQuery) {
        for (Iterator<Query> iterator = ((DisjunctionMaxQuery) query).iterator(); iterator.hasNext();) {
          extract(iterator.next(), terms);
        }
      } else if (query instanceof MultiTermQuery && expandMultiTermQuery) {
        MultiTermQuery mtq = ((MultiTermQuery)query);
        if(mtq.getRewriteMethod() != MultiTermQuery.SCORING_BOOLEAN_QUERY_REWRITE) {
          mtq = (MultiTermQuery) mtq.clone();
          mtq.setRewriteMethod(MultiTermQuery.SCORING_BOOLEAN_QUERY_REWRITE);
          query = mtq;
        }
        FakeReader fReader = new FakeReader();
        MultiTermQuery.SCORING_BOOLEAN_QUERY_REWRITE.rewrite(fReader, mtq);
        if (fReader.field != null) {
          IndexReader ir = getReaderForField(fReader.field);
          extract(query.rewrite(ir), terms);
        }
      } else if (query instanceof MultiPhraseQuery) {
        final MultiPhraseQuery mpq = (MultiPhraseQuery) query;
        final List<Term[]> termArrays = mpq.getTermArrays();
        final int[] positions = mpq.getPositions();
        if (positions.length > 0) {
  
          int maxPosition = positions[positions.length - 1];
          for (int i = 0; i < positions.length - 1; ++i) {
            if (positions[i] > maxPosition) {
              maxPosition = positions[i];
            }
          }
  
          @SuppressWarnings("unchecked")
          final List<SpanQuery>[] disjunctLists = new List[maxPosition + 1];
          int distinctPositions = 0;
  
          for (int i = 0; i < termArrays.size(); ++i) {
            final Term[] termArray = termArrays.get(i);
            List<SpanQuery> disjuncts = disjunctLists[positions[i]];
            if (disjuncts == null) {
              disjuncts = (disjunctLists[positions[i]] = new ArrayList<SpanQuery>(termArray.length));
              ++distinctPositions;
            }
            for (int j = 0; j < termArray.length; ++j) {
              disjuncts.add(new SpanTermQuery(termArray[j]));
            }
          }
  
          int positionGaps = 0;
          int position = 0;
          final SpanQuery[] clauses = new SpanQuery[distinctPositions];
          for (int i = 0; i < disjunctLists.length; ++i) {
            List<SpanQuery> disjuncts = disjunctLists[i];
            if (disjuncts != null) {
              clauses[position++] = new SpanOrQuery(disjuncts
                  .toArray(new SpanQuery[disjuncts.size()]));
            } else {
              ++positionGaps;
            }
          }
  
          final int slop = mpq.getSlop();
          final boolean inorder = (slop == 0);
  
          SpanNearQuery sp = new SpanNearQuery(clauses, slop + positionGaps, inorder);
          sp.setBoost(query.getBoost());
          extractWeightedSpanTerms(terms, sp);
        }
      }
    }
  
    /**
     * Fills a <code>Map</code> with <@link WeightedSpanTerm>s using the terms from the supplied <code>SpanQuery</code>.
     * 
     * @param terms
     *          Map to place created WeightedSpanTerms in
     * @param spanQuery
     *          SpanQuery to extract Terms from
     * @throws IOException
     */
    private void extractWeightedSpanTerms(Map<String,WeightedSpanTerm> terms, SpanQuery spanQuery) throws IOException {
      Set<String> fieldNames;
  
      if (fieldName == null) {
        fieldNames = new HashSet<String>();
        collectSpanQueryFields(spanQuery, fieldNames);
      } else {
        fieldNames = new HashSet<String>(1);
        fieldNames.add(fieldName);
      }
      // To support the use of the default field name
      if (defaultField != null) {
        fieldNames.add(defaultField);
      }
      
      Map<String, SpanQuery> queries = new HashMap<String, SpanQuery>();
   
      Set<Term> nonWeightedTerms = new HashSet<Term>();
      final boolean mustRewriteQuery = mustRewriteQuery(spanQuery);
      if (mustRewriteQuery) {
        for (final String field : fieldNames) {
          final SpanQuery rewrittenQuery = (SpanQuery) spanQuery.rewrite(getReaderForField(field));
          queries.put(field, rewrittenQuery);
          rewrittenQuery.extractTerms(nonWeightedTerms);
        }
      } else {
        spanQuery.extractTerms(nonWeightedTerms);
      }
  
      List<PositionSpan> spanPositions = new ArrayList<PositionSpan>();
  
      for (final String field : fieldNames) {
  
        IndexReader reader = getReaderForField(field);
        final Spans spans;
        if (mustRewriteQuery) {
          spans = queries.get(field).getSpans(reader);
        } else {
          spans = spanQuery.getSpans(reader);
        }
  
  
        // collect span positions
        while (spans.next()) {
          spanPositions.add(new PositionSpan(spans.start(), spans.end() - 1));
        }
        
      }
  
      if (spanPositions.size() == 0) {
        // no spans found
        return;
      }
  
      for (final Term queryTerm :  nonWeightedTerms) {
  
        if (fieldNameComparator(queryTerm.field())) {
          WeightedSpanTerm weightedSpanTerm = terms.get(queryTerm.text());
  
          if (weightedSpanTerm == null) {
            weightedSpanTerm = new WeightedSpanTerm(spanQuery.getBoost(), queryTerm.text());
            weightedSpanTerm.addPositionSpans(spanPositions);
            weightedSpanTerm.positionSensitive = true;
            terms.put(queryTerm.text(), weightedSpanTerm);
          } else {
            if (spanPositions.size() > 0) {
              weightedSpanTerm.addPositionSpans(spanPositions);
            }
          }
        }
      }
    }
  
    /**
     * Fills a <code>Map</code> with <@link WeightedSpanTerm>s using the terms from the supplied <code>Query</code>.
     * 
     * @param terms
     *          Map to place created WeightedSpanTerms in
     * @param query
     *          Query to extract Terms from
     * @throws IOException
     */
    private void extractWeightedTerms(Map<String,WeightedSpanTerm> terms, Query query) throws IOException {
      Set<Term> nonWeightedTerms = new HashSet<Term>();
      query.extractTerms(nonWeightedTerms);
  
      for (final Term queryTerm : nonWeightedTerms) {
  
        if (fieldNameComparator(queryTerm.field())) {
          WeightedSpanTerm weightedSpanTerm = new WeightedSpanTerm(query.getBoost(), queryTerm.text());
          terms.put(queryTerm.text(), weightedSpanTerm);
        }
      }
    }
  
    /**
     * Necessary to implement matches for queries against <code>defaultField</code>
     */
    private boolean fieldNameComparator(String fieldNameToCheck) {
      boolean rv = fieldName == null || fieldNameToCheck == fieldName
          || fieldNameToCheck == defaultField;
      return rv;
    }
  
    private IndexReader getReaderForField(String field) throws IOException {
      if(wrapToCaching && !cachedTokenStream && !(tokenStream instanceof CachingTokenFilter)) {
        tokenStream = new CachingTokenFilter(new OffsetLimitTokenFilter(tokenStream, maxDocCharsToAnalyze));
        cachedTokenStream = true;
      }
      IndexReader reader = readers.get(field);
      if (reader == null) {
        MemoryIndex indexer = new MemoryIndex();
        indexer.addField(field, new OffsetLimitTokenFilter(tokenStream, maxDocCharsToAnalyze));
        tokenStream.reset();
        IndexSearcher searcher = indexer.createSearcher();
        reader = searcher.getIndexReader();
        readers.put(field, reader);
      }
  
      return reader;
    }
  
    /**
     * Creates a Map of <code>WeightedSpanTerms</code> from the given <code>Query</code> and <code>TokenStream</code>.
     * 
     * <p>
     * 
     * @param query
     *          that caused hit
     * @param tokenStream
     *          of text to be highlighted
     * @return Map containing WeightedSpanTerms
     * @throws IOException
     */
    public Map<String,WeightedSpanTerm> getWeightedSpanTerms(Query query, TokenStream tokenStream)
        throws IOException {
      return getWeightedSpanTerms(query, tokenStream, null);
    }
  
    /**
     * Creates a Map of <code>WeightedSpanTerms</code> from the given <code>Query</code> and <code>TokenStream</code>.
     * 
     * <p>
     * 
     * @param query
     *          that caused hit
     * @param tokenStream
     *          of text to be highlighted
     * @param fieldName
     *          restricts Term's used based on field name
     * @return Map containing WeightedSpanTerms
     * @throws IOException
     */
    public Map<String,WeightedSpanTerm> getWeightedSpanTerms(Query query, TokenStream tokenStream,
        String fieldName) throws IOException {
      if (fieldName != null) {
        this.fieldName = StringHelper.intern(fieldName);
      } else {
        this.fieldName = null;
      }
  
      Map<String,WeightedSpanTerm> terms = new PositionCheckingMap<String>();
      this.tokenStream = tokenStream;
      try {
        extract(query, terms);
      } finally {
        closeReaders();
      }
  
      return terms;
    }
  
    /**
     * Creates a Map of <code>WeightedSpanTerms</code> from the given <code>Query</code> and <code>TokenStream</code>. Uses a supplied
     * <code>IndexReader</code> to properly weight terms (for gradient highlighting).
     * 
     * <p>
     * 
     * @param query
     *          that caused hit
     * @param tokenStream
     *          of text to be highlighted
     * @param fieldName
     *          restricts Term's used based on field name
     * @param reader
     *          to use for scoring
     * @return Map of WeightedSpanTerms with quasi tf/idf scores
     * @throws IOException
     */
    public Map<String,WeightedSpanTerm> getWeightedSpanTermsWithScores(Query query, TokenStream tokenStream, String fieldName,
        IndexReader reader) throws IOException {
      if (fieldName != null) {
        this.fieldName = StringHelper.intern(fieldName);
      } else {
        this.fieldName = null;
      }
      this.tokenStream = tokenStream;
  
      Map<String,WeightedSpanTerm> terms = new PositionCheckingMap<String>();
      extract(query, terms);
  
      int totalNumDocs = reader.numDocs();
      Set<String> weightedTerms = terms.keySet();
      Iterator<String> it = weightedTerms.iterator();
  
      try {
        while (it.hasNext()) {
          WeightedSpanTerm weightedSpanTerm = terms.get(it.next());
          int docFreq = reader.docFreq(new Term(fieldName, weightedSpanTerm.term));
          // docFreq counts deletes
          if(totalNumDocs < docFreq) {
            docFreq = totalNumDocs;
          }
          // IDF algorithm taken from DefaultSimilarity class
          float idf = (float) (Math.log((float) totalNumDocs / (double) (docFreq + 1)) + 1.0);
          weightedSpanTerm.weight *= idf;
        }
      } finally {
  
        closeReaders();
      }
  
      return terms;
    }
    
    private void collectSpanQueryFields(SpanQuery spanQuery, Set<String> fieldNames) {
      if (spanQuery instanceof FieldMaskingSpanQuery) {
        collectSpanQueryFields(((FieldMaskingSpanQuery)spanQuery).getMaskedQuery(), fieldNames);
      } else if (spanQuery instanceof SpanFirstQuery) {
        collectSpanQueryFields(((SpanFirstQuery)spanQuery).getMatch(), fieldNames);
      } else if (spanQuery instanceof SpanNearQuery) {
        for (final SpanQuery clause : ((SpanNearQuery)spanQuery).getClauses()) {
          collectSpanQueryFields(clause, fieldNames);
        }
      } else if (spanQuery instanceof SpanNotQuery) {
        collectSpanQueryFields(((SpanNotQuery)spanQuery).getInclude(), fieldNames);
      } else if (spanQuery instanceof SpanOrQuery) {
        for (final SpanQuery clause : ((SpanOrQuery)spanQuery).getClauses()) {
          collectSpanQueryFields(clause, fieldNames);
        }
      } else {
        fieldNames.add(spanQuery.getField());
      }
    }
    
    private boolean mustRewriteQuery(SpanQuery spanQuery) {
      if (!expandMultiTermQuery) {
        return false; // Will throw UnsupportedOperationException in case of a SpanRegexQuery.
      } else if (spanQuery instanceof FieldMaskingSpanQuery) {
        return mustRewriteQuery(((FieldMaskingSpanQuery)spanQuery).getMaskedQuery());
      } else if (spanQuery instanceof SpanFirstQuery) {
        return mustRewriteQuery(((SpanFirstQuery)spanQuery).getMatch());
      } else if (spanQuery instanceof SpanNearQuery) {
        for (final SpanQuery clause : ((SpanNearQuery)spanQuery).getClauses()) {
          if (mustRewriteQuery(clause)) {
            return true;
          }
        }
        return false; 
      } else if (spanQuery instanceof SpanNotQuery) {
        SpanNotQuery spanNotQuery = (SpanNotQuery)spanQuery;
        return mustRewriteQuery(spanNotQuery.getInclude()) || mustRewriteQuery(spanNotQuery.getExclude());
      } else if (spanQuery instanceof SpanOrQuery) {
        for (final SpanQuery clause : ((SpanOrQuery)spanQuery).getClauses()) {
          if (mustRewriteQuery(clause)) {
            return true;
          }
        }
        return false; 
      } else if (spanQuery instanceof SpanTermQuery) {
        return false;
      } else {
        return true;
      }
    }
    
    /**
     * This class makes sure that if both position sensitive and insensitive
     * versions of the same term are added, the position insensitive one wins.
     */
    static private class PositionCheckingMap<K> extends HashMap<K,WeightedSpanTerm> {
  
      @Override
      public void putAll(Map<? extends K,? extends WeightedSpanTerm> m) {
        for (Map.Entry<? extends K,? extends WeightedSpanTerm> entry : m.entrySet())
          this.put(entry.getKey(), entry.getValue());
      }
  
      @Override
      public WeightedSpanTerm put(K key, WeightedSpanTerm value) {
        WeightedSpanTerm prev = super.put(key, value);
        if (prev == null) return prev;
        WeightedSpanTerm prevTerm = prev;
        WeightedSpanTerm newTerm = value;
        if (!prevTerm.positionSensitive) {
          newTerm.positionSensitive = false;
        }
        return prev;
      }
      
    }
    
    public boolean getExpandMultiTermQuery() {
      return expandMultiTermQuery;
    }
  
    public void setExpandMultiTermQuery(boolean expandMultiTermQuery) {
      this.expandMultiTermQuery = expandMultiTermQuery;
    }
    
    public boolean isCachedTokenStream() {
      return cachedTokenStream;
    }
    
    public TokenStream getTokenStream() {
      return tokenStream;
    }
    
    /**
     * By default, {@link TokenStream}s that are not of the type
     * {@link CachingTokenFilter} are wrapped in a {@link CachingTokenFilter} to
     * ensure an efficient reset - if you are already using a different caching
     * {@link TokenStream} impl and you don't want it to be wrapped, set this to
     * false.
     * 
     * @param wrap
     */
    public void setWrapIfNotCachingTokenFilter(boolean wrap) {
      this.wrapToCaching = wrap;
    }
    
    /**
     * 
     * A fake IndexReader class to extract the field from a MultiTermQuery
     * 
     */
    static final class FakeReader extends FilterIndexReader {
  
      private static final IndexReader EMPTY_MEMORY_INDEX_READER =
        new MemoryIndex().createSearcher().getIndexReader();
      
      String field;
  
      FakeReader() {
        super(EMPTY_MEMORY_INDEX_READER);
      }
  
      @Override
      public TermEnum terms(final Term t) throws IOException {
        // only set first fieldname, maybe use a Set?
        if (t != null && field == null)
          field = t.field();
        return super.terms(t);
      }
  
  
    }
  
    protected final void setMaxDocCharsToAnalyze(int maxDocCharsToAnalyze) {
      this.maxDocCharsToAnalyze = maxDocCharsToAnalyze;
    }
  
  }
  