/*
 * 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.search;

import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.tests.util.LuceneTestCase;
import org.apache.lucene.tests.util.TestUtil;
import org.apache.lucene.util.BitDocIdSet;
import org.apache.lucene.util.FixedBitSet;

public class TestConjunctionDISI extends LuceneTestCase {

  private static TwoPhaseIterator approximation(
      DocIdSetIterator iterator, final FixedBitSet confirmed) {
    DocIdSetIterator approximation;
    if (random().nextBoolean()) {
      approximation = anonymizeIterator(iterator);
    } else {
      approximation = iterator;
    }
    return new TwoPhaseIterator(approximation) {

      @Override
      public boolean matches() {
        return confirmed.get(approximation.docID());
      }

      @Override
      public float matchCost() {
        return 5; // #operations in FixedBitSet#get()
      }
    };
  }

  /**
   * Return an anonym class so that ConjunctionDISI cannot optimize it like it does eg. for
   * BitSetIterators.
   */
  private static DocIdSetIterator anonymizeIterator(DocIdSetIterator it) {
    return new DocIdSetIterator() {

      @Override
      public int nextDoc() throws IOException {
        return it.nextDoc();
      }

      @Override
      public int docID() {
        return it.docID();
      }

      @Override
      public long cost() {
        return it.docID();
      }

      @Override
      public int advance(int target) throws IOException {
        return it.advance(target);
      }
    };
  }

  private static Scorer scorer(TwoPhaseIterator twoPhaseIterator) {
    return scorer(TwoPhaseIterator.asDocIdSetIterator(twoPhaseIterator), twoPhaseIterator);
  }

  private static class FakeWeight extends Weight {

    protected FakeWeight() {
      super(new MatchNoDocsQuery());
    }

    @Override
    public Explanation explain(LeafReaderContext context, int doc) throws IOException {
      return null;
    }

    @Override
    public Scorer scorer(LeafReaderContext context) throws IOException {
      return null;
    }

    @Override
    public boolean isCacheable(LeafReaderContext ctx) {
      return false;
    }
  }

  /**
   * Create a {@link Scorer} that wraps the given {@link DocIdSetIterator}. It also accepts a {@link
   * TwoPhaseIterator} view, which is exposed in {@link Scorer#twoPhaseIterator()}. When the
   * two-phase view is not null, then {@link DocIdSetIterator#nextDoc()} and {@link
   * DocIdSetIterator#advance(int)} will raise an exception in order to make sure that {@link
   * ConjunctionDISI} takes advantage of the {@link TwoPhaseIterator} view.
   */
  private static Scorer scorer(DocIdSetIterator it, TwoPhaseIterator twoPhaseIterator) {
    return new Scorer(new FakeWeight()) {

      @Override
      public DocIdSetIterator iterator() {
        return new DocIdSetIterator() {

          @Override
          public int docID() {
            return it.docID();
          }

          @Override
          public int nextDoc() throws IOException {
            if (twoPhaseIterator != null) {
              throw new UnsupportedOperationException(
                  "ConjunctionDISI should call the two-phase iterator");
            }
            return it.nextDoc();
          }

          @Override
          public int advance(int target) throws IOException {
            if (twoPhaseIterator != null) {
              throw new UnsupportedOperationException(
                  "ConjunctionDISI should call the two-phase iterator");
            }
            return it.advance(target);
          }

          @Override
          public long cost() {
            if (twoPhaseIterator != null) {
              throw new UnsupportedOperationException(
                  "ConjunctionDISI should call the two-phase iterator");
            }
            return it.cost();
          }
        };
      }

      @Override
      public TwoPhaseIterator twoPhaseIterator() {
        return twoPhaseIterator;
      }

      @Override
      public int docID() {
        if (twoPhaseIterator != null) {
          throw new UnsupportedOperationException(
              "ConjunctionDISI should call the two-phase iterator");
        }
        return it.docID();
      }

      @Override
      public float score() throws IOException {
        return 0;
      }

      @Override
      public float getMaxScore(int upTo) throws IOException {
        return 0;
      }
    };
  }

  private static FixedBitSet randomSet(int maxDoc) {
    final int step = TestUtil.nextInt(random(), 1, 10);
    FixedBitSet set = new FixedBitSet(maxDoc);
    for (int doc = random().nextInt(step);
        doc < maxDoc;
        doc += TestUtil.nextInt(random(), 1, step)) {
      set.set(doc);
    }
    return set;
  }

  private static FixedBitSet clearRandomBits(FixedBitSet other) {
    final FixedBitSet set = new FixedBitSet(other.length());
    set.or(other);
    for (int i = 0; i < set.length(); ++i) {
      if (random().nextBoolean()) {
        set.clear(i);
      }
    }
    return set;
  }

  private static FixedBitSet intersect(FixedBitSet[] bitSets) {
    final FixedBitSet intersection = new FixedBitSet(bitSets[0].length());
    intersection.or(bitSets[0]);
    for (int i = 1; i < bitSets.length; ++i) {
      intersection.and(bitSets[i]);
    }
    return intersection;
  }

  private static FixedBitSet toBitSet(int maxDoc, DocIdSetIterator iterator) throws IOException {
    final FixedBitSet set = new FixedBitSet(maxDoc);
    for (int doc = iterator.nextDoc(); doc != NO_MORE_DOCS; doc = iterator.nextDoc()) {
      set.set(doc);
    }
    return set;
  }

  // Test that the conjunction iterator is correct
  public void testConjunction() throws IOException {
    final int iters = atLeast(100);
    for (int iter = 0; iter < iters; ++iter) {
      final int maxDoc = TestUtil.nextInt(random(), 100, 10000);
      final int numIterators = TestUtil.nextInt(random(), 2, 5);
      final FixedBitSet[] sets = new FixedBitSet[numIterators];
      final Scorer[] iterators = new Scorer[numIterators];
      for (int i = 0; i < iterators.length; ++i) {
        final FixedBitSet set = randomSet(maxDoc);
        switch (random().nextInt(3)) {
          case 0:
            // simple iterator
            sets[i] = set;
            iterators[i] =
                new ConstantScoreScorer(
                    new FakeWeight(),
                    0f,
                    ScoreMode.TOP_SCORES,
                    anonymizeIterator(new BitDocIdSet(set).iterator()));
            break;
          case 1:
            // bitSet iterator
            sets[i] = set;
            iterators[i] =
                new ConstantScoreScorer(
                    new FakeWeight(), 0f, ScoreMode.TOP_SCORES, new BitDocIdSet(set).iterator());
            break;
          default:
            // scorer with approximation
            final FixedBitSet confirmed = clearRandomBits(set);
            sets[i] = confirmed;
            final TwoPhaseIterator approximation =
                approximation(new BitDocIdSet(set).iterator(), confirmed);
            iterators[i] = scorer(approximation);
            break;
        }
      }

      final DocIdSetIterator conjunction =
          ConjunctionUtils.intersectScorers(Arrays.asList(iterators));
      assertEquals(intersect(sets), toBitSet(maxDoc, conjunction));
    }
  }

  // Test that the conjunction approximation is correct
  public void testConjunctionApproximation() throws IOException {
    final int iters = atLeast(100);
    for (int iter = 0; iter < iters; ++iter) {
      final int maxDoc = TestUtil.nextInt(random(), 100, 10000);
      final int numIterators = TestUtil.nextInt(random(), 2, 5);
      final FixedBitSet[] sets = new FixedBitSet[numIterators];
      final Scorer[] iterators = new Scorer[numIterators];
      boolean hasApproximation = false;
      for (int i = 0; i < iterators.length; ++i) {
        final FixedBitSet set = randomSet(maxDoc);
        if (random().nextBoolean()) {
          // simple iterator
          sets[i] = set;
          iterators[i] =
              new ConstantScoreScorer(
                  new FakeWeight(),
                  0f,
                  ScoreMode.COMPLETE_NO_SCORES,
                  new BitDocIdSet(set).iterator());
        } else {
          // scorer with approximation
          final FixedBitSet confirmed = clearRandomBits(set);
          sets[i] = confirmed;
          final TwoPhaseIterator approximation =
              approximation(new BitDocIdSet(set).iterator(), confirmed);
          iterators[i] = scorer(approximation);
          hasApproximation = true;
        }
      }

      final DocIdSetIterator conjunction =
          ConjunctionUtils.intersectScorers(Arrays.asList(iterators));
      TwoPhaseIterator twoPhaseIterator = TwoPhaseIterator.unwrap(conjunction);
      assertEquals(hasApproximation, twoPhaseIterator != null);
      if (hasApproximation) {
        assertEquals(
            intersect(sets),
            toBitSet(maxDoc, TwoPhaseIterator.asDocIdSetIterator(twoPhaseIterator)));
      }
    }
  }

  // This test makes sure that when nesting scorers with ConjunctionDISI, confirmations are pushed
  // to the root.
  public void testRecursiveConjunctionApproximation() throws IOException {
    final int iters = atLeast(100);
    for (int iter = 0; iter < iters; ++iter) {
      final int maxDoc = TestUtil.nextInt(random(), 100, 10000);
      final int numIterators = TestUtil.nextInt(random(), 2, 5);
      final FixedBitSet[] sets = new FixedBitSet[numIterators];
      Scorer conjunction = null;
      boolean hasApproximation = false;
      for (int i = 0; i < numIterators; ++i) {
        final FixedBitSet set = randomSet(maxDoc);
        final Scorer newIterator;
        switch (random().nextInt(3)) {
          case 0:
            // simple iterator
            sets[i] = set;
            newIterator =
                new ConstantScoreScorer(
                    new FakeWeight(),
                    0f,
                    ScoreMode.TOP_SCORES,
                    anonymizeIterator(new BitDocIdSet(set).iterator()));
            break;
          case 1:
            // bitSet iterator
            sets[i] = set;
            newIterator =
                new ConstantScoreScorer(
                    new FakeWeight(), 0f, ScoreMode.TOP_SCORES, new BitDocIdSet(set).iterator());
            break;
          default:
            // scorer with approximation
            final FixedBitSet confirmed = clearRandomBits(set);
            sets[i] = confirmed;
            final TwoPhaseIterator approximation =
                approximation(new BitDocIdSet(set).iterator(), confirmed);
            newIterator = scorer(approximation);
            hasApproximation = true;
            break;
        }
        if (conjunction == null) {
          conjunction = newIterator;
        } else {
          final DocIdSetIterator conj =
              ConjunctionUtils.intersectScorers(Arrays.asList(conjunction, newIterator));
          conjunction = scorer(conj, TwoPhaseIterator.unwrap(conj));
        }
      }

      TwoPhaseIterator twoPhaseIterator = conjunction.twoPhaseIterator();
      assertEquals(hasApproximation, twoPhaseIterator != null);
      if (hasApproximation) {
        assertEquals(
            intersect(sets),
            toBitSet(maxDoc, TwoPhaseIterator.asDocIdSetIterator(twoPhaseIterator)));
      } else {
        assertEquals(intersect(sets), toBitSet(maxDoc, conjunction.iterator()));
      }
    }
  }

  public void testCollapseSubConjunctions(boolean wrapWithScorer) throws IOException {
    final int iters = atLeast(100);
    for (int iter = 0; iter < iters; ++iter) {
      final int maxDoc = TestUtil.nextInt(random(), 100, 10000);
      final int numIterators = TestUtil.nextInt(random(), 5, 10);
      final FixedBitSet[] sets = new FixedBitSet[numIterators];
      final List<Scorer> scorers = new LinkedList<>();
      for (int i = 0; i < numIterators; ++i) {
        final FixedBitSet set = randomSet(maxDoc);
        if (random().nextBoolean()) {
          // simple iterator
          sets[i] = set;
          scorers.add(
              new ConstantScoreScorer(
                  new FakeWeight(), 0f, ScoreMode.TOP_SCORES, new BitDocIdSet(set).iterator()));
        } else {
          // scorer with approximation
          final FixedBitSet confirmed = clearRandomBits(set);
          sets[i] = confirmed;
          final TwoPhaseIterator approximation =
              approximation(new BitDocIdSet(set).iterator(), confirmed);
          scorers.add(scorer(approximation));
        }
      }

      // make some sub sequences into sub conjunctions
      final int subIters = atLeast(3);
      for (int subIter = 0; subIter < subIters && scorers.size() > 3; ++subIter) {
        final int subSeqStart = TestUtil.nextInt(random(), 0, scorers.size() - 2);
        final int subSeqEnd = TestUtil.nextInt(random(), subSeqStart + 2, scorers.size());
        List<Scorer> subIterators = scorers.subList(subSeqStart, subSeqEnd);
        Scorer subConjunction;
        if (wrapWithScorer) {
          subConjunction =
              new ConjunctionScorer(new FakeWeight(), subIterators, Collections.emptyList());
        } else {
          subConjunction =
              new ConstantScoreScorer(
                  new FakeWeight(),
                  0f,
                  ScoreMode.TOP_SCORES,
                  ConjunctionUtils.intersectScorers(subIterators));
        }
        scorers.set(subSeqStart, subConjunction);
        int toRemove = subSeqEnd - subSeqStart - 1;
        while (toRemove-- > 0) {
          scorers.remove(subSeqStart + 1);
        }
      }
      if (scorers.size() == 1) {
        // ConjunctionDISI needs two iterators
        scorers.add(
            new ConstantScoreScorer(
                new FakeWeight(), 0f, ScoreMode.TOP_SCORES, DocIdSetIterator.all(maxDoc)));
      }

      final DocIdSetIterator conjunction = ConjunctionUtils.intersectScorers(scorers);
      assertEquals(intersect(sets), toBitSet(maxDoc, conjunction));
    }
  }

  public void testCollapseSubConjunctionDISIs() throws IOException {
    testCollapseSubConjunctions(false);
  }

  public void testCollapseSubConjunctionScorers() throws IOException {
    testCollapseSubConjunctions(true);
  }

  public void testIllegalAdvancementOfSubIteratorsTripsAssertion() throws IOException {
    assumeTrue("Assertions must be enabled for this test!", LuceneTestCase.assertsAreEnabled);
    int maxDoc = 100;
    final int numIterators = TestUtil.nextInt(random(), 2, 5);
    FixedBitSet set = randomSet(maxDoc);

    DocIdSetIterator[] iterators = new DocIdSetIterator[numIterators];
    for (int i = 0; i < iterators.length; ++i) {
      iterators[i] = new BitDocIdSet(set).iterator();
    }
    final DocIdSetIterator conjunction =
        ConjunctionUtils.intersectIterators(Arrays.asList(iterators));
    int idx = TestUtil.nextInt(random(), 0, iterators.length - 1);
    iterators[idx]
        .nextDoc(); // illegally advancing one of the sub-iterators outside of the conjunction
    // iterator
    AssertionError ex = expectThrows(AssertionError.class, () -> conjunction.nextDoc());
    assertEquals("Sub-iterators of ConjunctionDISI are not on the same document!", ex.getMessage());
  }

  public void testBitSetConjunctionDISIDocIDOnExhaust() throws IOException {
    int numBitSetIterators = TestUtil.nextInt(random(), 2, 5);
    DocIdSetIterator[] iterators = new DocIdSetIterator[numBitSetIterators + 1];

    // Create sparse DocIdSetIterator with a single match that is greater than lengths of bitset
    // iterators
    int maxBitSetLength = 1000;
    int minBitSetLength = 2;
    int leadMaxDoc = maxBitSetLength + 1;
    iterators[iterators.length - 1] = DocIdSetIterator.range(leadMaxDoc, leadMaxDoc + 1);

    for (int i = 0; i < numBitSetIterators; i++) {
      int bitSetLength = TestUtil.nextInt(random(), minBitSetLength, maxBitSetLength);
      FixedBitSet bitSet = new FixedBitSet(bitSetLength);
      bitSet.set(0, bitSetLength - 1);
      iterators[i] = new BitDocIdSet(bitSet).iterator();
    }

    final DocIdSetIterator conjunction =
        ConjunctionUtils.intersectIterators(Arrays.asList(iterators));

    assertEquals(NO_MORE_DOCS, conjunction.nextDoc());
    assertEquals(NO_MORE_DOCS, conjunction.docID());
  }
}
