package com.jeeww.lucene.demo;

/*
 * 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.Closeable;
import java.io.IOException;
import java.text.ParseException;

import org.apache.lucene.analysis.core.WhitespaceAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.DoubleField;
import org.apache.lucene.document.Field;
import org.apache.lucene.expressions.Expression;
import org.apache.lucene.expressions.SimpleBindings;
import org.apache.lucene.expressions.js.JavascriptCompiler;
import org.apache.lucene.facet.DrillDownQuery;
import org.apache.lucene.facet.DrillSideways;
import org.apache.lucene.facet.FacetResult;
import org.apache.lucene.facet.Facets;
import org.apache.lucene.facet.FacetsCollector;
import org.apache.lucene.facet.FacetsConfig;
import org.apache.lucene.facet.range.DoubleRange;
import org.apache.lucene.facet.range.DoubleRangeFacetCounts;
import org.apache.lucene.facet.taxonomy.TaxonomyReader;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.queries.BooleanFilter;
import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.Filter;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.NumericRangeFilter;
import org.apache.lucene.search.SortField;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.RAMDirectory;
import org.apache.lucene.util.SloppyMath;

/**
 * Shows simple usage of dynamic range faceting, using the expressions module to calculate distance.
 */
public class DistanceFacetsExample implements Closeable {

    final DoubleRange ONE_KM = new DoubleRange("< 1 km", 0.0, true, 1.0, false);
    final DoubleRange TWO_KM = new DoubleRange("< 2 km", 0.0, true, 2.0, false);
    final DoubleRange FIVE_KM = new DoubleRange("< 5 km", 0.0, true, 5.0, false);
    final DoubleRange TEN_KM = new DoubleRange("< 10 km", 0.0, true, 10.0, false);

    private final Directory indexDir = new RAMDirectory();
    private IndexSearcher searcher;
    private final FacetsConfig config = new FacetsConfig();

    /** The "home" latitude. */
    public final static double ORIGIN_LATITUDE = 40.7143528;

    /** The "home" longitude. */
    public final static double ORIGIN_LONGITUDE = -74.0059731;

    /**
     * Radius of the Earth in KM NOTE: this is approximate, because the earth is a bit wider at the equator than the
     * poles. See http://en.wikipedia.org/wiki/Earth_radius
     */
    public final static double EARTH_RADIUS_KM = 6371.01;

    /** Empty constructor */
    public DistanceFacetsExample() {
    }

    /** Build the example index. */
    public void index() throws IOException {
        IndexWriter writer = new IndexWriter(indexDir, new IndexWriterConfig(FacetExamples.EXAMPLES_VER,
                new WhitespaceAnalyzer(FacetExamples.EXAMPLES_VER)));

        // TODO: we could index in radians instead ... saves all the conversions in getBoundingBoxFilter

        // Add documents with latitude/longitude location:
        Document doc = new Document();
        doc.add(new DoubleField("latitude", 40.759011, Field.Store.NO));
        doc.add(new DoubleField("longitude", -73.9844722, Field.Store.NO));
        writer.addDocument(doc);

        doc = new Document();
        doc.add(new DoubleField("latitude", 40.718266, Field.Store.NO));
        doc.add(new DoubleField("longitude", -74.007819, Field.Store.NO));
        writer.addDocument(doc);

        doc = new Document();
        doc.add(new DoubleField("latitude", 40.7051157, Field.Store.NO));
        doc.add(new DoubleField("longitude", -74.0088305, Field.Store.NO));
        writer.addDocument(doc);

        // Open near-real-time searcher
        searcher = new IndexSearcher(DirectoryReader.open(writer, true));
        writer.close();
    }

    private ValueSource getDistanceValueSource() {
        Expression distance;
        try {
            distance = JavascriptCompiler.compile("haversin(" + ORIGIN_LATITUDE + "," + ORIGIN_LONGITUDE
                    + ",latitude,longitude)");
        } catch (ParseException pe) {
            // Should not happen
            throw new RuntimeException(pe);
        }
        SimpleBindings bindings = new SimpleBindings();
        bindings.add(new SortField("latitude", SortField.Type.DOUBLE));
        bindings.add(new SortField("longitude", SortField.Type.DOUBLE));

        return distance.getValueSource(bindings);
    }

    /**
     * Given a latitude and longitude (in degrees) and the maximum great circle (surface of the earth) distance, returns
     * a simple Filter bounding box to "fast match" candidates.
     */
    public static Filter getBoundingBoxFilter(double originLat, double originLng, double maxDistanceKM) {

        // Basic bounding box geo math from
        // http://JanMatuschek.de/LatitudeLongitudeBoundingCoordinates,
        // licensed under creative commons 3.0:
        // http://creativecommons.org/licenses/by/3.0

        // TODO: maybe switch to recursive prefix tree instead
        // (in lucene/spatial)? It should be more efficient
        // since it's a 2D trie...

        // Degrees -> Radians:
        double originLatRadians = Math.toRadians(originLat);
        double originLngRadians = Math.toRadians(originLng);

        double angle = maxDistanceKM / (SloppyMath.earthDiameter(originLat) / 2.0);

        double minLat = originLatRadians - angle;
        double maxLat = originLatRadians + angle;

        double minLng;
        double maxLng;
        if (minLat > Math.toRadians(-90) && maxLat < Math.toRadians(90)) {
            double delta = Math.asin(Math.sin(angle) / Math.cos(originLatRadians));
            minLng = originLngRadians - delta;
            if (minLng < Math.toRadians(-180)) {
                minLng += 2 * Math.PI;
            }
            maxLng = originLngRadians + delta;
            if (maxLng > Math.toRadians(180)) {
                maxLng -= 2 * Math.PI;
            }
        } else {
            // The query includes a pole!
            minLat = Math.max(minLat, Math.toRadians(-90));
            maxLat = Math.min(maxLat, Math.toRadians(90));
            minLng = Math.toRadians(-180);
            maxLng = Math.toRadians(180);
        }

        BooleanFilter f = new BooleanFilter();

        // Add latitude range filter:
        f.add(NumericRangeFilter.newDoubleRange("latitude", Math.toDegrees(minLat), Math.toDegrees(maxLat), true, true),
                BooleanClause.Occur.MUST);

        // Add longitude range filter:
        if (minLng > maxLng) {
            // The bounding box crosses the international date
            // line:
            BooleanFilter lonF = new BooleanFilter();
            lonF.add(NumericRangeFilter.newDoubleRange("longitude", Math.toDegrees(minLng), null, true, true),
                    BooleanClause.Occur.SHOULD);
            lonF.add(NumericRangeFilter.newDoubleRange("longitude", null, Math.toDegrees(maxLng), true, true),
                    BooleanClause.Occur.SHOULD);
            f.add(lonF, BooleanClause.Occur.MUST);
        } else {
            f.add(NumericRangeFilter.newDoubleRange("longitude", Math.toDegrees(minLng), Math.toDegrees(maxLng), true,
                    true), BooleanClause.Occur.MUST);
        }

        return f;
    }

    /** User runs a query and counts facets. */
    public FacetResult search() throws IOException {

        FacetsCollector fc = new FacetsCollector();

        searcher.search(new MatchAllDocsQuery(), fc);

        Facets facets = new DoubleRangeFacetCounts("field", getDistanceValueSource(), fc, getBoundingBoxFilter(
                ORIGIN_LATITUDE, ORIGIN_LONGITUDE, 10.0), ONE_KM, TWO_KM, FIVE_KM, TEN_KM);

        return facets.getTopChildren(10, "field");
    }

    /** User drills down on the specified range. */
    public TopDocs drillDown(DoubleRange range) throws IOException {

        // Passing no baseQuery means we drill down on all
        // documents ("browse only"):
        DrillDownQuery q = new DrillDownQuery(null);
        final ValueSource vs = getDistanceValueSource();
        q.add("field", range.getFilter(getBoundingBoxFilter(ORIGIN_LATITUDE, ORIGIN_LONGITUDE, range.max), vs));
        DrillSideways ds = new DrillSideways(searcher, config, (TaxonomyReader) null) {
            @Override
            protected Facets buildFacetsResult(FacetsCollector drillDowns, FacetsCollector[] drillSideways,
                    String[] drillSidewaysDims) throws IOException {
                assert drillSideways.length == 1;
                return new DoubleRangeFacetCounts("field", vs, drillSideways[0], ONE_KM, TWO_KM, FIVE_KM, TEN_KM);
            }
        };
        return ds.search(q, 10).hits;
    }

    @Override
    public void close() throws IOException {
        searcher.getIndexReader().close();
        indexDir.close();
    }

    /** Runs the search and drill-down examples and prints the results. */
    @SuppressWarnings("unchecked")
    public static void main(String[] args) throws Exception {
        DistanceFacetsExample example = new DistanceFacetsExample();
        example.index();

        System.out.println("Distance facet counting example:");
        System.out.println("-----------------------");
        System.out.println(example.search());

        System.out.println("\n");
        System.out.println("Distance facet drill-down example (field/< 2 km):");
        System.out.println("---------------------------------------------");
        TopDocs hits = example.drillDown(example.TWO_KM);
        System.out.println(hits.totalHits + " totalHits");

        example.close();
    }
}
