package org.apache.lucene.facet.util;

import java.io.IOException;
import java.util.Random;

import org.apache.lucene.analysis.MockAnalyzer;
import org.apache.lucene.analysis.MockTokenizer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.Field.Index;
import org.apache.lucene.document.Field.Store;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.RandomIndexWriter;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.DocIdSet;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TermQuery;
import org.apache.lucene.store.Directory;
import org.apache.lucene.util.OpenBitSet;
import org.apache.lucene.util.OpenBitSetDISI;
import org.junit.Test;

import org.apache.lucene.util.LuceneTestCase;
import org.apache.lucene.facet.search.ScoredDocIDs;
import org.apache.lucene.facet.search.ScoredDocIDsIterator;
import org.apache.lucene.facet.search.ScoredDocIdCollector;

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

public class TestScoredDocIDsUtils extends LuceneTestCase {

  @Test
  public void testComplementIterator() throws Exception {
    final int n = atLeast(10000);
    final OpenBitSet bits = new OpenBitSet(n);
    for (int i = 0; i < 5 * n; i++) {
      bits.flip(random.nextInt(n));
    }
    
    OpenBitSet verify = new OpenBitSet(n);
    verify.or(bits);

    ScoredDocIDs scoredDocIDs = ScoredDocIdsUtils.createScoredDocIds(bits, n); 

    Directory dir = newDirectory();
    IndexReader reader = createReaderWithNDocs(random, n, dir);
    try { 
      assertEquals(n - verify.cardinality(), ScoredDocIdsUtils.getComplementSet(scoredDocIDs, 
        reader).size());
    } finally {
      reader.close();
      dir.close();
    }
  }

  @Test
  public void testAllDocs() throws Exception {
    int maxDoc = 3;
    Directory dir = newDirectory();
    IndexReader reader = createReaderWithNDocs(random, maxDoc, dir);
    try {
      ScoredDocIDs all = ScoredDocIdsUtils.createAllDocsScoredDocIDs(reader);
      assertEquals("invalid size", maxDoc, all.size());
      ScoredDocIDsIterator iter = all.iterator();
      int doc = 0;
      while (iter.next()) {
        assertEquals("invalid doc ID: " + iter.getDocID(), doc++, iter.getDocID());
        assertEquals("invalid score: " + iter.getScore(), ScoredDocIDsIterator.DEFAULT_SCORE, iter.getScore(), 0.0f);
      }
      assertEquals("invalid maxDoc: " + doc, maxDoc, doc);
      
      DocIdSet docIDs = all.getDocIDs();
      assertTrue("should be cacheable", docIDs.isCacheable());
      DocIdSetIterator docIDsIter = docIDs.iterator();
      assertEquals("nextDoc() hasn't been called yet", -1, docIDsIter.docID());
      assertEquals(0, docIDsIter.nextDoc());
      assertEquals(1, docIDsIter.advance(1));
      // if advance is smaller than current doc, advance to cur+1.
      assertEquals(2, docIDsIter.advance(0));
    } finally {
      reader.close();
      dir.close();
    }
  }
  
  @Test
  public void testWithDeletions() throws Exception {
    int N_DOCS = 100;

    DocumentFactory docFactory = new DocumentFactory(N_DOCS) {
      @Override
      public boolean markedDeleted(int docNum) {
        return (docNum % 3 == 0 ||        // every 3rd documents, including first 
            docNum == numDocs - 1 ||     // last document
            docNum == numDocs / 2 ||     // 3 consecutive documents in the middle
            docNum == 1 + numDocs / 2 ||
            docNum == 2 + numDocs / 2);
      }
      
      // every 6th document (starting from the 2nd) would contain 'alpha'
      @Override
      public boolean haveAlpha(int docNum) {
        return (docNum % 6 == 1);
      }
    };
    
    Directory dir = newDirectory();
    IndexReader reader = createReaderWithNDocs(random, N_DOCS, docFactory, dir);
    try {
      ScoredDocIDs allDocs = ScoredDocIdsUtils.createAllDocsScoredDocIDs(reader);
      ScoredDocIDsIterator it = allDocs.iterator();
      int numIteratedDocs = 0;
      while (it.next()) {
        numIteratedDocs++;
        int docNum = it.getDocID();
        assertNull(
            "Deleted docs must not appear in the allDocsScoredDocIds set: " + docNum,
            reader.document(docNum).getFieldable("del"));
      }

      assertEquals("Wrong number of (live) documents", allDocs.size(), numIteratedDocs);

      // Get all 'alpha' documents
      ScoredDocIdCollector collector = ScoredDocIdCollector.create(reader.maxDoc(), false);
      Query q = new TermQuery(new Term(DocumentFactory.field, DocumentFactory.alphaTxt));
      IndexSearcher searcher = newSearcher(reader);
      searcher.search(q, collector);
      searcher.close();

      ScoredDocIDs scoredDocIds = collector.getScoredDocIDs();
      OpenBitSet resultSet = new OpenBitSetDISI(scoredDocIds.getDocIDs().iterator(), reader.maxDoc());
      
      // Getting the complement set of the query result
      ScoredDocIDs complementSet = ScoredDocIdsUtils.getComplementSet(scoredDocIds, reader);

      assertEquals("Number of documents in complement set mismatch",
          reader.numDocs() - scoredDocIds.size(), complementSet.size());

      // now make sure the documents in the complement set are not deleted
      // and not in the original result set
      ScoredDocIDsIterator compIterator = complementSet.iterator();
      while (compIterator.next()) {
        int docNum = compIterator.getDocID();
        assertFalse(
            "Complement-Set must not contain deleted documents (doc="+docNum+")",
            reader.isDeleted(docNum));
        assertNull(
            "Complement-Set must not contain docs from the original set (doc="+ docNum+")",
            reader.document(docNum).getFieldable("del"));
        assertFalse(
            "Complement-Set must not contain docs from the original set (doc="+docNum+")",
            resultSet.fastGet(docNum));
      }
    } finally {
      reader.close();
      dir.close();
    }
  }
  
  /**
   * Creates an index with n documents, this method is meant for testing purposes ONLY
   */
  static IndexReader createReaderWithNDocs(Random random, int nDocs, Directory directory) throws IOException {
    return createReaderWithNDocs(random, nDocs, new DocumentFactory(nDocs), directory);
  }

  private static class DocumentFactory {
    protected final static String field = "content";
    protected final static String delTxt = "delete";
    protected final static String alphaTxt = "alpha";
    
    private final static Field deletionMark = new Field(field, delTxt, Store.NO, Index.NOT_ANALYZED_NO_NORMS);
    private final static Field alphaContent = new Field(field, alphaTxt, Store.NO, Index.NOT_ANALYZED_NO_NORMS);
    
    protected final int numDocs;
    
    public DocumentFactory(int totalNumDocs) {
      this.numDocs = totalNumDocs;
    }
    
    public boolean markedDeleted(int docNum) {
      return false;
    }

    public Document getDoc(int docNum) {
      Document doc = new Document();
      if (markedDeleted(docNum)) {
        doc.add(deletionMark);
        // Add a special field for docs that are marked for deletion. Later we
        // assert that those docs are not returned by all-scored-doc-IDs.
        doc.add(new Field("del", Integer.toString(docNum), Store.YES, Index.NO));
      }

      if (haveAlpha(docNum)) {
        doc.add(alphaContent);
      }
      return doc;
    }

    public boolean haveAlpha(int docNum) {
      return false;
    }
  }

  static IndexReader createReaderWithNDocs(Random random, int nDocs, DocumentFactory docFactory, Directory dir) throws IOException {
    RandomIndexWriter writer = new RandomIndexWriter(random, dir,
        newIndexWriterConfig(random, TEST_VERSION_CURRENT,
            new MockAnalyzer(random, MockTokenizer.KEYWORD, false)));
    for (int docNum = 0; docNum < nDocs; docNum++) {
      writer.addDocument(docFactory.getDoc(docNum));
    }
    writer.close();

    // Delete documents marked for deletion
    IndexReader reader = IndexReader.open(dir, false);
    reader.deleteDocuments(new Term(DocumentFactory.field, DocumentFactory.delTxt));
    reader.close();

    // Open a fresh read-only reader with the deletions in place
    return IndexReader.open(dir, true);
  }
}
