package au.gov.immi.ngbs.test;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

import org.apache.lucene.analysis.core.WhitespaceAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field.Store;
import org.apache.lucene.document.IntField;
import org.apache.lucene.facet.index.FacetFields;
import org.apache.lucene.facet.params.FacetSearchParams;
import org.apache.lucene.facet.search.Aggregator;
import org.apache.lucene.facet.search.FacetArrays;
import org.apache.lucene.facet.search.FacetRequest;
import org.apache.lucene.facet.search.FacetResult;
import org.apache.lucene.facet.search.FacetResultNode;
import org.apache.lucene.facet.search.FacetsCollector;
import org.apache.lucene.facet.taxonomy.CategoryPath;
import org.apache.lucene.facet.taxonomy.TaxonomyReader;
import org.apache.lucene.facet.taxonomy.directory.DirectoryTaxonomyReader;
import org.apache.lucene.facet.taxonomy.directory.DirectoryTaxonomyWriter;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.RAMDirectory;
import org.apache.lucene.util.IntsRef;
import org.apache.lucene.util.Version;

public class SumFacetsTestCase {
	private Directory indexDir;
	private Directory taxoDir;
	private IndexWriter indexWriter;
	private DirectoryTaxonomyWriter taxoWriter;
	private FacetFields facetFields;
	private Map<String, Integer> indexedSum = new TreeMap<String, Integer>();

	public static void main(String[] args) throws IOException 
	{
		new SumFacetsTestCase();
	}
	
	public SumFacetsTestCase() throws IOException
	{
		indexDir = new RAMDirectory();
		taxoDir = new RAMDirectory();
		
	    indexWriter = new IndexWriter(indexDir, new IndexWriterConfig(Version.LUCENE_43, 
	            new WhitespaceAnalyzer(Version.LUCENE_43)));
        taxoWriter = new DirectoryTaxonomyWriter(taxoDir);
        facetFields = new FacetFields(taxoWriter);
		
		index("level1/level2A/level3A", 2);
		index("level1/level2B/level3A", 3);
		index("level1/level2A/level3A", 4);
		index("level1/level2A/level3B", 5);
		index("level1/level2B/level3A", 4);
		index("level1/level2B/level3B", 3);
		index("level1/level2C/level3A", 0);
		index("level1/level2A/level3A", 1);
		index("level1/level2A/level3A", 0);
		
	    indexWriter.close();
	    taxoWriter.close();
		
		System.out.println(indexedSum.toString());
		System.out.println(search().toString());
	}
	
	private void index(String facet, Integer alertCount) throws IOException
	{
		Document doc = new Document();
		
		List<CategoryPath> paths = new ArrayList<CategoryPath>();
		paths.add(new CategoryPath(facet, '/'));
        
        doc.add(new IntField("alertCount", alertCount, Store.YES));
        
        facetFields.addFields(doc, paths);
        
        indexWriter.addDocument(doc);
        
        indexWriter.commit();
        taxoWriter.commit();
        
        if (indexedSum.containsKey(facet))
        {
        	indexedSum.put(facet, indexedSum.get(facet) + alertCount);
        }
        else
        {
        	indexedSum.put(facet, alertCount);
        }
	}
	
	private Map<String, Integer> search() throws IOException
	{
        DirectoryReader indexReader = DirectoryReader.open(indexDir);
        IndexSearcher searcher = new IndexSearcher(indexReader);

        TaxonomyReader taxoReader = new DirectoryTaxonomyReader(taxoDir);

        FacetRequest facetRequest = new SumFacetRequest(new CategoryPath("level1"), Integer.MAX_VALUE, searcher, "alertCount");
        facetRequest.setDepth(4);
        
        FacetSearchParams fsp = new FacetSearchParams(facetRequest);
        
        FacetsCollector facetsColl = FacetsCollector.create(fsp, searcher.getIndexReader(), taxoReader);
        
        searcher.search(new MatchAllDocsQuery(), facetsColl);
        List<FacetResult> facetResults = facetsColl.getFacetResults();
        
        indexReader.close();
        taxoReader.close();
    
        Map<String, Integer> extractedResults = new TreeMap<String, Integer>();
        readFacetTree(facetResults.get(0).getFacetResultNode(), extractedResults);
        
        return extractedResults;
	}
	
	// flattens out the facet tree into a Map
    private void readFacetTree(FacetResultNode node, Map<String, Integer> map)
    {
        if (node.subResults.size() == 0)
        {
            // we are at the bottom of the tree
            String path = getFacetPath(node.label.components);
            map.put(path, (int)node.value);
        }
        else
        {
            for (FacetResultNode subNode : node.subResults)
            {
                readFacetTree(subNode, map);
            }
        }
    }

    private String getFacetPath(String[] components)
    {
        StringBuilder builder = new StringBuilder();
        for (int i = 0; i < components.length; i++)
        {
            if (i > 0)
            {
                builder.append("/");
            }
            builder.append(components[i]);
        }
        
        return builder.toString();
    }
    
    private class SumFacetRequest extends FacetRequest
    {
        private IndexSearcher searcher;
        private String fieldName;
        
        public SumFacetRequest(CategoryPath path, int numResults, IndexSearcher searcher, String fieldName)
        {
            super(path, numResults);
            this.searcher = searcher;
            this.fieldName = fieldName;
        }

        @Override
        public Aggregator createAggregator(boolean useComplements, FacetArrays arrays, TaxonomyReader taxonomy) {
          int[] a = arrays.getIntArray();
          if (useComplements) {
            throw new UnsupportedOperationException("SumFacetRequest cannot use complements");
          }
          return new SummingAggregator(a, searcher, fieldName);
        }

        @Override
        public FacetArraysSource getFacetArraysSource()
        {
            return FacetArraysSource.INT;
        }

        @Override
        public double getValueOf(FacetArrays arrays, int ordinal)
        {
            return arrays.getIntArray()[ordinal];
        }
    }
    
    private class SummingAggregator implements Aggregator
    {
        protected int[] sumArray;
        private IndexSearcher searcher;
        private String fieldName;

        public SummingAggregator(int[] sumArray, IndexSearcher searcher, String fieldName)
        {
            this.sumArray = sumArray;
            this.searcher = searcher;
            this.fieldName = fieldName;
        }

        public void aggregate(int docID, float score, IntsRef ordinals) throws IOException
        {
            Document doc = searcher.doc(docID);
            int value = doc.getField(fieldName).numericValue().intValue();
            
            for (int i = 0; i < ordinals.length; i++)
            {
                sumArray[ordinals.ints[i]] += value;
            }
        }

        @Override
        public boolean equals(Object obj)
        {
            if (obj == null || obj.getClass() != this.getClass())
            {
                return false;
            }
            SummingAggregator that = (SummingAggregator) obj;
            return that.sumArray == this.sumArray;
        }

        @Override
        public int hashCode()
        {
            return sumArray == null ? 0 : sumArray.hashCode();
        }

        public boolean setNextReader(AtomicReaderContext context) throws IOException
        {
            return true;
        }
    }

}