package org.apache.lucene.search;

/*
 * 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.Comparator;
import java.util.List;

import org.apache.lucene.util.ArrayUtil;

/**
 * A Scorer for OR like queries, counterpart of <code>ConjunctionScorer</code>.
 * This Scorer implements {@link Scorer#advance(int)} and uses advance() on the given Scorers.
 * 
 * This implementation uses the minimumMatch constraint actively to efficiently
 * prune the number of candidates, it is hence a mixture between a pure DisjunctionScorer
 * and a ConjunctionScorer.
 */
class MinShouldMatchSumScorer extends Scorer {

  /** The overall number of non-finalized scorers */
  private int numScorers;
  /** The minimum number of scorers that should match */
  private final int mm;

  /** A static array of all subscorers sorted by decreasing cost */
  private final Scorer sortedSubScorers[];
  /** A monotonically increasing index into the array pointing to the next subscorer that is to be excluded */
  private int sortedSubScorersIdx = 0;

  private final Scorer subScorers[]; // the first numScorers-(mm-1) entries are valid
  private int nrInHeap; // 0..(numScorers-(mm-1)-1)

  /** mmStack is supposed to contain the most costly subScorers that still did
   *  not run out of docs, sorted by increasing sparsity of docs returned by that subScorer.
   *  For now, the cost of subscorers is assumed to be inversely correlated with sparsity.
   */
  private final Scorer mmStack[]; // of size mm-1: 0..mm-2, always full

  /** The document number of the current match. */
  private int doc = -1;
  /** The number of subscorers that provide the current match. */
  protected int nrMatchers = -1;
  private double score = Float.NaN;

  /**
   * Construct a <code>MinShouldMatchSumScorer</code>.
   * 
   * @param weight The weight to be used.
   * @param subScorers A collection of at least two subscorers.
   * @param minimumNrMatchers The positive minimum number of subscorers that should
   * match to match this query.
   * <br>When <code>minimumNrMatchers</code> is bigger than
   * the number of <code>subScorers</code>, no matches will be produced.
   * <br>When minimumNrMatchers equals the number of subScorers,
   * it is more efficient to use <code>ConjunctionScorer</code>.
   */
  public MinShouldMatchSumScorer(Weight weight, List<Scorer> subScorers, int minimumNrMatchers) throws IOException {
    super(weight);
    this.nrInHeap = this.numScorers = subScorers.size();

    if (minimumNrMatchers <= 0) {
      throw new IllegalArgumentException("Minimum nr of matchers must be positive");
    }
    if (numScorers <= 1) {
      throw new IllegalArgumentException("There must be at least 2 subScorers");
    }

    this.mm = minimumNrMatchers;
    this.sortedSubScorers = subScorers.toArray(new Scorer[this.numScorers]);
    // sorting by decreasing subscorer cost should be inversely correlated with
    // next docid (assuming costs are due to generating many postings)
    ArrayUtil.timSort(sortedSubScorers, new Comparator<Scorer>() {
      @Override
      public int compare(Scorer o1, Scorer o2) {
        return Long.signum(o2.cost() - o1.cost());
      }
    });
    // take mm-1 most costly subscorers aside
    this.mmStack = new Scorer[mm-1];
    for (int i = 0; i < mm-1; i++) {
      mmStack[i] = sortedSubScorers[i];
    }
    nrInHeap -= mm-1;
    this.sortedSubScorersIdx = mm-1;
    // take remaining into heap, if any, and heapify
    this.subScorers = new Scorer[nrInHeap];
    for (int i = 0; i < nrInHeap; i++) {
      this.subScorers[i] = this.sortedSubScorers[mm-1+i];
    }
    minheapHeapify();
    assert minheapCheck();
  }
  
  /**
   * Construct a <code>DisjunctionScorer</code>, using one as the minimum number
   * of matching subscorers.
   */
  public MinShouldMatchSumScorer(Weight weight, List<Scorer> subScorers) throws IOException {
    this(weight, subScorers, 1);
  }

  @Override
  public final Collection<ChildScorer> getChildren() {
    ArrayList<ChildScorer> children = new ArrayList<>(numScorers);
    for (int i = 0; i < numScorers; i++) {
      children.add(new ChildScorer(subScorers[i], "SHOULD"));
    }
    return children;
  }
  
  @Override
  public int nextDoc() throws IOException {
    assert doc != NO_MORE_DOCS;
    while (true) {
      // to remove current doc, call next() on all subScorers on current doc within heap
      while (subScorers[0].docID() == doc) {
        if (subScorers[0].nextDoc() != NO_MORE_DOCS) {
          minheapSiftDown(0);
        } else {
          minheapRemoveRoot();
          numScorers--;
          if (numScorers < mm) {
            return doc = NO_MORE_DOCS;
          }
        }
        //assert minheapCheck();
      }

      evaluateSmallestDocInHeap();

      if (nrMatchers >= mm) { // doc satisfies mm constraint
        break;
      }
    }
    return doc;
  }
  
  private void evaluateSmallestDocInHeap() throws IOException {
    // within heap, subScorer[0] now contains the next candidate doc
    doc = subScorers[0].docID();
    if (doc == NO_MORE_DOCS) {
      nrMatchers = Integer.MAX_VALUE; // stop looping
      return;
    }
    // 1. score and count number of matching subScorers within heap
    score = subScorers[0].score();
    nrMatchers = 1;
    countMatches(1);
    countMatches(2);
    // 2. score and count number of matching subScorers within stack,
    // short-circuit: stop when mm can't be reached for current doc, then perform on heap next()
    // TODO instead advance() might be possible, but complicates things
    for (int i = mm-2; i >= 0; i--) { // first advance sparsest subScorer
      if (mmStack[i].docID() >= doc || mmStack[i].advance(doc) != NO_MORE_DOCS) {
        if (mmStack[i].docID() == doc) { // either it was already on doc, or got there via advance()
          nrMatchers++;
          score += mmStack[i].score();
        } else { // scorer advanced to next after doc, check if enough scorers left for current doc
          if (nrMatchers + i < mm) { // too few subScorers left, abort advancing
            return; // continue looping TODO consider advance() here
          }
        }
      } else { // subScorer exhausted
        numScorers--;
        if (numScorers < mm) { // too few subScorers left
          doc = NO_MORE_DOCS;
          nrMatchers = Integer.MAX_VALUE; // stop looping
          return;
        }
        if (mm-2-i > 0) {
          // shift RHS of array left
          System.arraycopy(mmStack, i+1, mmStack, i, mm-2-i);
        }
        // find next most costly subScorer within heap TODO can this be done better?
        while (!minheapRemove(sortedSubScorers[sortedSubScorersIdx++])) {
          //assert minheapCheck();
        }
        // add the subScorer removed from heap to stack
        mmStack[mm-2] = sortedSubScorers[sortedSubScorersIdx-1];
        
        if (nrMatchers + i < mm) { // too few subScorers left, abort advancing
          return; // continue looping TODO consider advance() here
        }
      }
    }
  }

  // TODO: this currently scores, but so did the previous impl
  // TODO: remove recursion.
  // TODO: consider separating scoring out of here, then modify this
  // and afterNext() to terminate when nrMatchers == minimumNrMatchers
  // then also change freq() to just always compute it from scratch
  private void countMatches(int root) throws IOException {
    if (root < nrInHeap && subScorers[root].docID() == doc) {
      nrMatchers++;
      score += subScorers[root].score();
      countMatches((root<<1)+1);
      countMatches((root<<1)+2);
    }
  }

  /**
   * Returns the score of the current document matching the query. Initially
   * invalid, until {@link #nextDoc()} is called the first time.
   */
  @Override
  public float score() throws IOException {
    return (float) score;
  }

  @Override
  public int docID() {
    return doc;
  }

  @Override
  public int freq() throws IOException {
    return nrMatchers;
  }

  /**
   * Advances to the first match beyond the current whose document number is
   * greater than or equal to a given target. <br>
   * The implementation uses the advance() method on the subscorers.
   * 
   * @param target the target document number.
   * @return the document whose number is greater than or equal to the given
   *         target, or -1 if none exist.
   */
  @Override
  public int advance(int target) throws IOException {
    if (numScorers < mm)
      return doc = NO_MORE_DOCS;
    // advance all Scorers in heap at smaller docs to at least target
    while (subScorers[0].docID() < target) {
      if (subScorers[0].advance(target) != NO_MORE_DOCS) {
        minheapSiftDown(0);
      } else {
        minheapRemoveRoot();
        numScorers--;
        if (numScorers < mm) {
          return doc = NO_MORE_DOCS;
        }
      }
      //assert minheapCheck();
    }

    evaluateSmallestDocInHeap();

    if (nrMatchers >= mm) {
      return doc;
    } else {
      return nextDoc();
    }
  }
  
  @Override
  public long cost() {
    // cost for merging of lists analog to DisjunctionSumScorer
    long costCandidateGeneration = 0;
    for (int i = 0; i < nrInHeap; i++)
      costCandidateGeneration += subScorers[i].cost();
    // TODO is cost for advance() different to cost for iteration + heap merge
    //      and how do they compare overall to pure disjunctions? 
    final float c1 = 1.0f,
                c2 = 1.0f; // maybe a constant, maybe a proportion between costCandidateGeneration and sum(subScorer_to_be_advanced.cost())?
    return (long) (
           c1 * costCandidateGeneration +        // heap-merge cost
           c2 * costCandidateGeneration * (mm-1) // advance() cost
           );
  }
  
  /**
   * Organize subScorers into a min heap with scorers generating the earliest document on top.
   */
  protected final void minheapHeapify() {
    for (int i = (nrInHeap >> 1) - 1; i >= 0; i--) {
      minheapSiftDown(i);
    }
  }
  
  /**
   * The subtree of subScorers at root is a min heap except possibly for its root element.
   * Bubble the root down as required to make the subtree a heap.
   */
  protected final void minheapSiftDown(int root) {
    // TODO could this implementation also move rather than swapping neighbours?
    Scorer scorer = subScorers[root];
    int doc = scorer.docID();
    int i = root;
    while (i <= (nrInHeap >> 1) - 1) {
      int lchild = (i << 1) + 1;
      Scorer lscorer = subScorers[lchild];
      int ldoc = lscorer.docID();
      int rdoc = Integer.MAX_VALUE, rchild = (i << 1) + 2;
      Scorer rscorer = null;
      if (rchild < nrInHeap) {
        rscorer = subScorers[rchild];
        rdoc = rscorer.docID();
      }
      if (ldoc < doc) {
        if (rdoc < ldoc) {
          subScorers[i] = rscorer;
          subScorers[rchild] = scorer;
          i = rchild;
        } else {
          subScorers[i] = lscorer;
          subScorers[lchild] = scorer;
          i = lchild;
        }
      } else if (rdoc < doc) {
        subScorers[i] = rscorer;
        subScorers[rchild] = scorer;
        i = rchild;
      } else {
        return;
      }
    }
  }

  protected final void minheapSiftUp(int i) {
    Scorer scorer = subScorers[i];
    final int doc = scorer.docID();
    // find right place for scorer
    while (i > 0) {
      int parent = (i - 1) >> 1;
      Scorer pscorer = subScorers[parent];
      int pdoc = pscorer.docID();
      if (pdoc > doc) { // move root down, make space
        subScorers[i] = subScorers[parent];
        i = parent;
      } else { // done, found right place
        break;
      }
    }
    subScorers[i] = scorer;
  }

  /**
   * Remove the root Scorer from subScorers and re-establish it as a heap
   */
  protected final void minheapRemoveRoot() {
    if (nrInHeap == 1) {
      //subScorers[0] = null; // not necessary
      nrInHeap = 0;
    } else {
      nrInHeap--;
      subScorers[0] = subScorers[nrInHeap];
      //subScorers[nrInHeap] = null; // not necessary
      minheapSiftDown(0);
    }
  }
  
  /**
   * Removes a given Scorer from the heap by placing end of heap at that
   * position and bubbling it either up or down
   */
  protected final boolean minheapRemove(Scorer scorer) {
    // find scorer: O(nrInHeap)
    for (int i = 0; i < nrInHeap; i++) {
      if (subScorers[i] == scorer) { // remove scorer
        subScorers[i] = subScorers[--nrInHeap];
        //if (i != nrInHeap) subScorers[nrInHeap] = null; // not necessary
        minheapSiftUp(i);
        minheapSiftDown(i);
        return true;
      }
    }
    return false; // scorer already exhausted
  }
  
  boolean minheapCheck() {
    return minheapCheck(0);
  }
  private boolean minheapCheck(int root) {
    if (root >= nrInHeap)
      return true;
    int lchild = (root << 1) + 1;
    int rchild = (root << 1) + 2;
    if (lchild < nrInHeap && subScorers[root].docID() > subScorers[lchild].docID())
      return false;
    if (rchild < nrInHeap && subScorers[root].docID() > subScorers[rchild].docID())
      return false;
    return minheapCheck(lchild) && minheapCheck(rchild);
  }
  
}