/*
 * Copyright 2017 Organizations participating in ISAAC, ISAAC's KOMET, and SOLOR development include the 
         US Veterans Health Administration, OSHERA, and the Health Services Platform Consortium..
 *
 * Licensed 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 sh.isaac.provider.query.lucene;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.function.Predicate;

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.LeafCollector;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TopScoreDocCollector;

/**
 *
 * @author kec
 * @author <a href="mailto:daniel.armbrust.list@gmail.com">Dan Armbrust</a>
 */
public class IsaacFilteredCollectorManager implements CollectorManager<IsaacFilteredCollectorManager.IsaacFilteredCollector, TopDocs> {

   final Predicate<Integer> filter;
   final int sizeLimit;
   final ScoreDoc after;
   final ArrayList<IsaacFilteredCollector> collectors = new ArrayList<>();
   protected static final Logger LOG = LogManager.getLogger();

   /**
    * 
    * @param filter - optional - arbitrary filtering
    * @param sizeLimit - max number of results to return
    * @param after - optional - specify to begin on a future page of results.  Should be the last ScoreDoc from a previous query.
    */
   public IsaacFilteredCollectorManager(Predicate<Integer> filter, int sizeLimit, ScoreDoc after) {
      this.filter = filter;
      this.sizeLimit = sizeLimit;
      this.after = after;
   }
   
   TopDocs getTopDocs() throws IOException {
      return reduce(collectors);
   }

   @Override
   public IsaacFilteredCollector newCollector() throws IOException {
      IsaacFilteredCollector collector =  new IsaacFilteredCollector();
      collectors.add(collector);
      return collector;
   }

   @Override
   public TopDocs reduce(Collection<IsaacFilteredCollectorManager.IsaacFilteredCollector> collectors) throws IOException {
      final TopDocs[] topDocs = new TopDocs[collectors.size()];
      int i = 0;
      int totalSize = 0;
      for (IsaacFilteredCollectorManager.IsaacFilteredCollector collector : collectors) {
         topDocs[i] = collector.topDocs();
         totalSize = totalSize + topDocs[i++].scoreDocs.length;
      }
      return TopDocs.merge(0, Math.min(sizeLimit, totalSize), topDocs, true);
   }

   public class IsaacFilteredCollector implements Collector {
      HashMap<LeafReaderContext, LeafCollector> leafDelegates = new HashMap<>();
      TopScoreDocCollector delegateCollector = TopScoreDocCollector.create(sizeLimit, after);
      
      @Override
      public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
         if (leafDelegates.get(context) == null)
         {
            leafDelegates.put(context, new LeafCollector() 
            {
               LeafCollector internalDelegate = delegateCollector.getLeafCollector(context);
               
               @Override
               public void setScorer(Scorer scorer) throws IOException {
                  internalDelegate.setScorer(scorer);
                  
               }
               
               @Override
               public void collect(int doc) throws IOException {
                  boolean filterPass = false;
                  if (filter != null)
                  {
                     //This little optimization prevents us from re-evaluating filters 
                     int afterDoc = (after == null ? 0 : after.doc) - context.docBase;
                     if (doc < afterDoc)
                     {
                        //TODO [DAN 1] test this optimization, it differs slightly from what they do in Lucene (where they calc score first)
                        //but I can't see that is necessary...
                         LOG.trace("skipping filter eval for item before page 1");
                         filterPass = false;
                     }
                     else
                     {
                        final Document document = context.reader().document(doc);
                        final int componentNid = document.getField(LuceneIndexer.FIELD_COMPONENT_NID)
                                .numericValue()
                                .intValue();
                        filterPass = filter.test(componentNid);
                     }
                  }
                  else
                  {
                     filterPass = true;
                  }
                  
                  if (filterPass)
                  {
                     internalDelegate.collect(doc);
                  }
               }
            });
         }
         return leafDelegates.get(context);
      }

      @Override
      public boolean needsScores() {
         return delegateCollector.needsScores();
      }
      
      public TopDocs topDocs() {
         return delegateCollector.topDocs();
      }
   }
}

@@ -4141,157 +4141,8 @@
 se;%0A
-                  if (filter != null)%0A                  %7B%0A                     //This little optimization prevents us from re-evaluating filters %0A   
     
@@ -4235,27 +4235,24 @@
             
-   
 if (doc %3C af
@@ -4269,39 +4269,34 @@
                 
-   %7B%0A  
+%7B%0A
                 
@@ -4304,192 +4304,126 @@
     
- 
 //T
-ODO %5BDAN 1%5D test this optimization, it differs slightly from what they do in Lucene (where they calc score first)%0A                        //but I can't see that is necessary...%0A   
+his little optimization prevents us from re-evaluating filters and from re-collecting the same docs more than once%0A
     
@@ -4472,16 +4472,28 @@
 er eval 
+and collect 
 for item
@@ -4524,35 +4524,32 @@
                 
-   
 filterPass = fal
@@ -4562,35 +4562,32 @@
                 
-  
- 
 %7D%0A              
@@ -4586,35 +4586,44 @@
             
-   else%0A   
+if (filter != null)%0A
             
@@ -4622,36 +4622,32 @@
               %7B%0A
-    
                 
@@ -4724,28 +4724,24 @@
             
-    
 final int co
@@ -4834,20 +4834,16 @@
         
-    
 .numeric
@@ -4878,20 +4878,16 @@
         
-    
 .intValu
@@ -4903,36 +4903,32 @@
                 
- 
-   
 filterPass = fil
@@ -4955,31 +4955,8 @@
 d);%0A
-                     %7D%0A
     

fix a bug where we were collecting too many search results on second (and subsequent pages) of results