/*
 * 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.
 */
package org.apache.lucene.facet;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;

import org.apache.lucene.facet.sortedset.SortedSetDocValuesFacetCounts;
import org.apache.lucene.facet.sortedset.SortedSetDocValuesFacetField;
import org.apache.lucene.facet.sortedset.SortedSetDocValuesReaderState;
import org.apache.lucene.facet.taxonomy.FastTaxonomyFacetCounts;
import org.apache.lucene.facet.taxonomy.TaxonomyReader;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.FilterCollector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.MultiCollector;
import org.apache.lucene.search.MultiCollectorManager;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TopFieldCollector;
import org.apache.lucene.search.TopFieldDocs;
import org.apache.lucene.search.TopScoreDocCollector;
import org.apache.lucene.util.ThreadInterruptedException;

Computes drill down and sideways counts for the provided DrillDownQuery. Drill sideways counts include alternative values/aggregates for the drill-down dimensions so that a dimension does not disappear after the user drills down into it. 
 Use one of the static search methods to do the search, and then get the hits and facet results from the returned DrillSidewaysResult. 
NOTE: this allocates one FacetsCollector for each drill-down, plus one. If your index has high number of facet labels then this will multiply your memory usage. 
@lucene.experimental
/**
 * Computes drill down and sideways counts for the provided
 * {@link DrillDownQuery}.  Drill sideways counts include
 * alternative values/aggregates for the drill-down
 * dimensions so that a dimension does not disappear after
 * the user drills down into it.
 * <p> Use one of the static search
 * methods to do the search, and then get the hits and facet
 * results from the returned {@link DrillSidewaysResult}.
 * <p><b>NOTE</b>: this allocates one {@link
 * FacetsCollector} for each drill-down, plus one.  If your
 * index has high number of facet labels then this will
 * multiply your memory usage.
 *
 * @lucene.experimental
 */
public class DrillSideways {

  
IndexSearcher passed to constructor. /**
   * {@link IndexSearcher} passed to constructor.
   */
  protected final IndexSearcher searcher;

  
TaxonomyReader passed to constructor. /**
   * {@link TaxonomyReader} passed to constructor.
   */
  protected final TaxonomyReader taxoReader;

  
SortedSetDocValuesReaderState passed to constructor; can be null. /**
   * {@link SortedSetDocValuesReaderState} passed to
   * constructor; can be null.
   */
  protected final SortedSetDocValuesReaderState state;

  
FacetsConfig passed to constructor. /**
   * {@link FacetsConfig} passed to constructor.
   */
  protected final FacetsConfig config;

  // These are only used for multi-threaded search
  private final ExecutorService executor;

  
Create a new DrillSideways instance. /**
   * Create a new {@code DrillSideways} instance.
   */
  public DrillSideways(IndexSearcher searcher, FacetsConfig config, TaxonomyReader taxoReader) {
    this(searcher, config, taxoReader, null);
  }

  
Create a new DrillSideways instance, assuming the categories were indexed with SortedSetDocValuesFacetField. /**
   * Create a new {@code DrillSideways} instance, assuming the categories were
   * indexed with {@link SortedSetDocValuesFacetField}.
   */
  public DrillSideways(IndexSearcher searcher, FacetsConfig config, SortedSetDocValuesReaderState state) {
    this(searcher, config, null, state);
  }

  
Create a new DrillSideways instance, where some dimensions were indexed with SortedSetDocValuesFacetField and others were indexed with FacetField. /**
   * Create a new {@code DrillSideways} instance, where some
   * dimensions were indexed with {@link
   * SortedSetDocValuesFacetField} and others were indexed
   * with {@link FacetField}.
   */
  public DrillSideways(IndexSearcher searcher, FacetsConfig config, TaxonomyReader taxoReader,
          SortedSetDocValuesReaderState state) {
    this(searcher, config, taxoReader, state, null);
  }

  
Create a new DrillSideways instance, where some dimensions were indexed with SortedSetDocValuesFacetField and others were indexed with FacetField. 

Use this constructor to use the concurrent implementation and/or the CollectorManager
/**
   * Create a new {@code DrillSideways} instance, where some
   * dimensions were indexed with {@link
   * SortedSetDocValuesFacetField} and others were indexed
   * with {@link FacetField}.
   * <p>
   * Use this constructor to use the concurrent implementation and/or the CollectorManager
   */
  public DrillSideways(IndexSearcher searcher, FacetsConfig config, TaxonomyReader taxoReader,
          SortedSetDocValuesReaderState state, ExecutorService executor) {
    this.searcher = searcher;
    this.config = config;
    this.taxoReader = taxoReader;
    this.state = state;
    this.executor = executor;
  }

  
Subclass can override to customize per-dim Facets
impl.
/**
   * Subclass can override to customize per-dim Facets
   * impl.
   */
  protected Facets buildFacetsResult(FacetsCollector drillDowns, FacetsCollector[] drillSideways,
          String[] drillSidewaysDims) throws IOException {

    Facets drillDownFacets;
    Map<String, Facets> drillSidewaysFacets = new HashMap<>();

    if (taxoReader != null) {
      drillDownFacets = new FastTaxonomyFacetCounts(taxoReader, config, drillDowns);
      if (drillSideways != null) {
        for (int i = 0; i < drillSideways.length; i++) {
          drillSidewaysFacets.put(drillSidewaysDims[i],
                  new FastTaxonomyFacetCounts(taxoReader, config, drillSideways[i]));
        }
      }
    } else {
      drillDownFacets = new SortedSetDocValuesFacetCounts(state, drillDowns);
      if (drillSideways != null) {
        for (int i = 0; i < drillSideways.length; i++) {
          drillSidewaysFacets.put(drillSidewaysDims[i], new SortedSetDocValuesFacetCounts(state, drillSideways[i]));
        }
      }
    }

    if (drillSidewaysFacets.isEmpty()) {
      return drillDownFacets;
    } else {
      return new MultiFacets(drillSidewaysFacets, drillDownFacets);
    }
  }

  
Search, collecting hits with a Collector, and computing drill down and sideways counts. /**
   * Search, collecting hits with a {@link Collector}, and
   * computing drill down and sideways counts.
   */
  public DrillSidewaysResult search(DrillDownQuery query, Collector hitCollector) throws IOException {

    Map<String, Integer> drillDownDims = query.getDims();

    FacetsCollector drillDownCollector = new FacetsCollector();

    if (drillDownDims.isEmpty()) {
      // There are no drill-down dims, so there is no
      // drill-sideways to compute:
      searcher.search(query, MultiCollector.wrap(hitCollector, drillDownCollector));
      return new DrillSidewaysResult(buildFacetsResult(drillDownCollector, null, null), null);
    }

    Query baseQuery = query.getBaseQuery();
    if (baseQuery == null) {
      // TODO: we could optimize this pure-browse case by
      // making a custom scorer instead:
      baseQuery = new MatchAllDocsQuery();
    }
    Query[] drillDownQueries = query.getDrillDownQueries();

    FacetsCollector[] drillSidewaysCollectors = new FacetsCollector[drillDownDims.size()];
    for (int i = 0; i < drillSidewaysCollectors.length; i++) {
      drillSidewaysCollectors[i] = new FacetsCollector();
    }

    DrillSidewaysQuery dsq =
            new DrillSidewaysQuery(baseQuery, drillDownCollector, drillSidewaysCollectors, drillDownQueries,
                    scoreSubDocsAtOnce());
    if (hitCollector.scoreMode().needsScores() == false) {
      // this is a horrible hack in order to make sure IndexSearcher will not
      // attempt to cache the DrillSidewaysQuery
      hitCollector = new FilterCollector(hitCollector) {
        @Override
        public ScoreMode scoreMode() {
          return ScoreMode.COMPLETE;
        }
      };
    }
    searcher.search(dsq, hitCollector);

    return new DrillSidewaysResult(buildFacetsResult(drillDownCollector, drillSidewaysCollectors,
            drillDownDims.keySet().toArray(new String[drillDownDims.size()])), null);
  }

  
Search, sorting by Sort, and computing drill down and sideways counts. /**
   * Search, sorting by {@link Sort}, and computing
   * drill down and sideways counts.
   */
  public DrillSidewaysResult search(DrillDownQuery query, Query filter, FieldDoc after, int topN, Sort sort,
          boolean doDocScores) throws IOException {
    if (filter != null) {
      query = new DrillDownQuery(config, filter, query);
    }
    if (sort != null) {
      int limit = searcher.getIndexReader().maxDoc();
      if (limit == 0) {
        limit = 1; // the collector does not alow numHits = 0
      }
      final int fTopN = Math.min(topN, limit);

      if (executor != null) { // We have an executor, let use the multi-threaded version

        final CollectorManager<TopFieldCollector, TopFieldDocs> collectorManager =
                new CollectorManager<TopFieldCollector, TopFieldDocs>() {

                  @Override
                  public TopFieldCollector newCollector() throws IOException {
                    return TopFieldCollector.create(sort, fTopN, after, Integer.MAX_VALUE);
                  }

                  @Override
                  public TopFieldDocs reduce(Collection<TopFieldCollector> collectors) throws IOException {
                    final TopFieldDocs[] topFieldDocs = new TopFieldDocs[collectors.size()];
                    int pos = 0;
                    for (TopFieldCollector collector : collectors)
                      topFieldDocs[pos++] = collector.topDocs();
                    return TopDocs.merge(sort, topN, topFieldDocs);
                  }

                };
        ConcurrentDrillSidewaysResult<TopFieldDocs> r = search(query, collectorManager);
        TopFieldDocs topDocs = r.collectorResult;
        if (doDocScores) {
          TopFieldCollector.populateScores(topDocs.scoreDocs, searcher, query);
        }
        return new DrillSidewaysResult(r.facets, topDocs);

      } else {

        final TopFieldCollector hitCollector =
                TopFieldCollector.create(sort, fTopN, after, Integer.MAX_VALUE);
        DrillSidewaysResult r = search(query, hitCollector);
        TopFieldDocs topDocs = hitCollector.topDocs();
        if (doDocScores) {
          TopFieldCollector.populateScores(topDocs.scoreDocs, searcher, query);
        }
        return new DrillSidewaysResult(r.facets, topDocs);
      }
    } else {
      return search(after, query, topN);
    }
  }

  
Search, sorting by score, and computing
drill down and sideways counts.
/**
   * Search, sorting by score, and computing
   * drill down and sideways counts.
   */
  public DrillSidewaysResult search(DrillDownQuery query, int topN) throws IOException {
    return search(null, query, topN);
  }

  
Search, sorting by score, and computing
drill down and sideways counts.
/**
   * Search, sorting by score, and computing
   * drill down and sideways counts.
   */
  public DrillSidewaysResult search(ScoreDoc after, DrillDownQuery query, int topN) throws IOException {
    int limit = searcher.getIndexReader().maxDoc();
    if (limit == 0) {
      limit = 1; // the collector does not alow numHits = 0
    }
    final int fTopN = Math.min(topN, limit);

    if (executor != null) {  // We have an executor, let use the multi-threaded version

      final CollectorManager<TopScoreDocCollector, TopDocs> collectorManager =
              new CollectorManager<TopScoreDocCollector, TopDocs>() {

                @Override
                public TopScoreDocCollector newCollector() throws IOException {
                  return TopScoreDocCollector.create(fTopN, after, Integer.MAX_VALUE);
                }

                @Override
                public TopDocs reduce(Collection<TopScoreDocCollector> collectors) throws IOException {
                  final TopDocs[] topDocs = new TopDocs[collectors.size()];
                  int pos = 0;
                  for (TopScoreDocCollector collector : collectors)
                    topDocs[pos++] = collector.topDocs();
                  return TopDocs.merge(topN, topDocs);
                }

              };
      ConcurrentDrillSidewaysResult<TopDocs> r = search(query, collectorManager);
      return new DrillSidewaysResult(r.facets, r.collectorResult);

    } else {

      TopScoreDocCollector hitCollector = TopScoreDocCollector.create(topN, after, Integer.MAX_VALUE);
      DrillSidewaysResult r = search(query, hitCollector);
      return new DrillSidewaysResult(r.facets, hitCollector.topDocs());
    }
  }

  
Override this and return true if your collector (e.g., ToParentBlockJoinCollector) expects all sub-scorers to be positioned on the document being collected. This will cause some performance loss; default is false. /**
   * Override this and return true if your collector
   * (e.g., {@code ToParentBlockJoinCollector}) expects all
   * sub-scorers to be positioned on the document being
   * collected.  This will cause some performance loss;
   * default is false.
   */
  protected boolean scoreSubDocsAtOnce() {
    return false;
  }

  
Result of a drill sideways search, including the Facets and TopDocs. /**
   * Result of a drill sideways search, including the
   * {@link Facets} and {@link TopDocs}.
   */
  public static class DrillSidewaysResult {
    
Combined drill down and sideways results.
/**
     * Combined drill down and sideways results.
     */
    public final Facets facets;

    
Hits.
/**
     * Hits.
     */
    public final TopDocs hits;

    
Sole constructor.
/**
     * Sole constructor.
     */
    public DrillSidewaysResult(Facets facets, TopDocs hits) {
      this.facets = facets;
      this.hits = hits;
    }
  }

  private static class CallableCollector implements Callable<CallableResult> {

    private final int pos;
    private final IndexSearcher searcher;
    private final Query query;
    private final CollectorManager<?, ?> collectorManager;

    private CallableCollector(int pos, IndexSearcher searcher, Query query, CollectorManager<?, ?> collectorManager) {
      this.pos = pos;
      this.searcher = searcher;
      this.query = query;
      this.collectorManager = collectorManager;
    }

    @Override
    public CallableResult call() throws Exception {
      return new CallableResult(pos, searcher.search(query, collectorManager));
    }
  }

  private static class CallableResult {

    private final int pos;
    private final Object result;

    private CallableResult(int pos, Object result) {
      this.pos = pos;
      this.result = result;
    }
  }

  private DrillDownQuery getDrillDownQuery(final DrillDownQuery query, Query[] queries,
          final String excludedDimension) {
    final DrillDownQuery ddl = new DrillDownQuery(config, query.getBaseQuery());
    query.getDims().forEach((dim, pos) -> {
      if (!dim.equals(excludedDimension))
        ddl.add(dim, queries[pos]);
    });
    return ddl.getDims().size() == queries.length ? null : ddl;
  }

  
Runs a search, using a CollectorManager to gather and merge search results /** Runs a search, using a {@link CollectorManager} to gather and merge search results */
  public <R> ConcurrentDrillSidewaysResult<R> search(final DrillDownQuery query,
          final CollectorManager<?, R> hitCollectorManager) throws IOException {

    final Map<String, Integer> drillDownDims = query.getDims();
    final List<CallableCollector> callableCollectors = new ArrayList<>(drillDownDims.size() + 1);

    // Add the main DrillDownQuery
    callableCollectors.add(new CallableCollector(-1, searcher, query,
            new MultiCollectorManager(new FacetsCollectorManager(), hitCollectorManager)));
    int i = 0;
    final Query[] filters = query.getDrillDownQueries();
    for (String dim : drillDownDims.keySet())
      callableCollectors.add(new CallableCollector(i++, searcher, getDrillDownQuery(query, filters, dim),
              new FacetsCollectorManager()));

    final FacetsCollector mainFacetsCollector;
    final FacetsCollector[] facetsCollectors = new FacetsCollector[drillDownDims.size()];
    final R collectorResult;

    try {
      // Run the query pool
      final List<Future<CallableResult>> futures = executor.invokeAll(callableCollectors);

      // Extract the results
      final Object[] mainResults = (Object[]) futures.get(0).get().result;
      mainFacetsCollector = (FacetsCollector) mainResults[0];
      collectorResult = (R) mainResults[1];
      for (i = 1; i < futures.size(); i++) {
        final CallableResult result = futures.get(i).get();
        facetsCollectors[result.pos] = (FacetsCollector) result.result;
      }
      // Fill the null results with the mainFacetsCollector
      for (i = 0; i < facetsCollectors.length; i++)
        if (facetsCollectors[i] == null)
          facetsCollectors[i] = mainFacetsCollector;

    } catch (InterruptedException e) {
      throw new ThreadInterruptedException(e);
    } catch (ExecutionException e) {
      throw new RuntimeException(e);
    }

    // build the facets and return the result
    return new ConcurrentDrillSidewaysResult<>(buildFacetsResult(mainFacetsCollector, facetsCollectors,
            drillDownDims.keySet().toArray(new String[drillDownDims.size()])), null, collectorResult);
  }

  
Result of a concurrent drill sideways search, including the Facets and TopDocs. /**
   * Result of a concurrent drill sideways search, including the
   * {@link Facets} and {@link TopDocs}.
   */
  public static class ConcurrentDrillSidewaysResult<R> extends DrillSidewaysResult {

    
The merged search results /** The merged search results */
    public final R collectorResult;

    
Sole constructor.
/**
     * Sole constructor.
     */
    ConcurrentDrillSidewaysResult(Facets facets, TopDocs hits, R collectorResult) {
      super(facets, hits);
      this.collectorResult = collectorResult;
    }
  }
}