/*
 * 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.benchmark.byTask.feeds;


import java.util.AbstractMap;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
import java.util.Set;

import org.locationtech.spatial4j.context.SpatialContext;
import org.locationtech.spatial4j.context.SpatialContextFactory;
import org.locationtech.spatial4j.shape.Point;
import org.locationtech.spatial4j.shape.Shape;
import org.apache.lucene.benchmark.byTask.utils.Config;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.spatial.SpatialStrategy;
import org.apache.lucene.spatial.composite.CompositeSpatialStrategy;
import org.apache.lucene.spatial.prefix.RecursivePrefixTreeStrategy;
import org.apache.lucene.spatial.prefix.tree.PackedQuadPrefixTree;
import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTreeFactory;
import org.apache.lucene.spatial.serialized.SerializedDVStrategy;

/**
 * Indexes spatial data according to a configured {@link SpatialStrategy} with optional
 * shape transformation via a configured {@link ShapeConverter}. The converter can turn points into
 * circles and bounding boxes, in order to vary the type of indexing performance tests.
 * Unless it's subclass-ed to do otherwise, this class configures a {@link SpatialContext},
 * {@link SpatialPrefixTree}, and {@link RecursivePrefixTreeStrategy}. The Strategy is made
 * available to a query maker via the static method {@link #getSpatialStrategy(int)}.
 * See spatial.alg for a listing of spatial parameters, in particular those starting with "spatial."
 * and "doc.spatial".
 */
public class SpatialDocMaker extends DocMaker {

  public static final String SPATIAL_FIELD = "spatial";

  //cache spatialStrategy by round number
  private static Map<Integer,SpatialStrategy> spatialStrategyCache = new HashMap<>();

  private SpatialStrategy strategy;
  private ShapeConverter shapeConverter;

  /**
   * Looks up the SpatialStrategy from the given round --
   * {@link org.apache.lucene.benchmark.byTask.utils.Config#getRoundNumber()}. It's an error
   * if it wasn't created already for this round -- when SpatialDocMaker is initialized.
   */
  public static SpatialStrategy getSpatialStrategy(int roundNumber) {
    SpatialStrategy result = spatialStrategyCache.get(roundNumber);
    if (result == null) {
      throw new IllegalStateException("Strategy should have been init'ed by SpatialDocMaker by now");
    }
    return result;
  }

  /**
   * Builds a SpatialStrategy from configuration options.
   */
  protected SpatialStrategy makeSpatialStrategy(final Config config) {
    //A Map view of Config that prefixes keys with "spatial."
    Map<String, String> configMap = new AbstractMap<String, String>() {
      @Override
      public Set<Entry<String, String>> entrySet() {
        throw new UnsupportedOperationException();
      }

      @Override
      public String get(Object key) {
        return config.get("spatial." + key, null);
      }
    };

    SpatialContext ctx = SpatialContextFactory.makeSpatialContext(configMap, null);

    return makeSpatialStrategy(config, configMap, ctx);
  }

  protected SpatialStrategy makeSpatialStrategy(final Config config, Map<String, String> configMap,
                                                SpatialContext ctx) {
    //TODO once strategies have factories, we could use them here.
    final String strategyName = config.get("spatial.strategy", "rpt");
    switch (strategyName) {
      case "rpt": return makeRPTStrategy(SPATIAL_FIELD, config, configMap, ctx);
      case "composite": return makeCompositeStrategy(config, configMap, ctx);
      //TODO add more as-needed
      default: throw new IllegalStateException("Unknown spatial.strategy: " + strategyName);
    }
  }

  protected RecursivePrefixTreeStrategy makeRPTStrategy(String spatialField, Config config,
                                                        Map<String, String> configMap, SpatialContext ctx) {
    //A factory for the prefix tree grid
    SpatialPrefixTree grid = SpatialPrefixTreeFactory.makeSPT(configMap, null, ctx);

    RecursivePrefixTreeStrategy strategy = new RecursivePrefixTreeStrategy(grid, spatialField);
    strategy.setPointsOnly(config.get("spatial.docPointsOnly", false));
    final boolean pruneLeafyBranches = config.get("spatial.pruneLeafyBranches", true);
    if (grid instanceof PackedQuadPrefixTree) {
      ((PackedQuadPrefixTree) grid).setPruneLeafyBranches(pruneLeafyBranches);
      strategy.setPruneLeafyBranches(false);//always leave it to packed grid, even though it isn't the same
    } else {
      strategy.setPruneLeafyBranches(pruneLeafyBranches);
    }

    int prefixGridScanLevel = config.get("query.spatial.prefixGridScanLevel", -4);
    if (prefixGridScanLevel < 0)
      prefixGridScanLevel = grid.getMaxLevels() + prefixGridScanLevel;
    strategy.setPrefixGridScanLevel(prefixGridScanLevel);

    double distErrPct = config.get("spatial.distErrPct", .025);//doc & query; a default
    strategy.setDistErrPct(distErrPct);
    return strategy;
  }

  protected SerializedDVStrategy makeSerializedDVStrategy(String spatialField, Config config,
                                                          Map<String, String> configMap, SpatialContext ctx) {
    return new SerializedDVStrategy(ctx, spatialField);
  }

  protected SpatialStrategy makeCompositeStrategy(Config config, Map<String, String> configMap, SpatialContext ctx) {
    final CompositeSpatialStrategy strategy = new CompositeSpatialStrategy(
        SPATIAL_FIELD, makeRPTStrategy(SPATIAL_FIELD + "_rpt", config, configMap, ctx),
        makeSerializedDVStrategy(SPATIAL_FIELD + "_sdv", config, configMap, ctx)
    );
    strategy.setOptimizePredicates(config.get("query.spatial.composite.optimizePredicates", true));
    return strategy;
  }

  @Override
  public void setConfig(Config config, ContentSource source) {
    super.setConfig(config, source);
    SpatialStrategy existing = spatialStrategyCache.get(config.getRoundNumber());
    if (existing == null) {
      //new round; we need to re-initialize
      strategy = makeSpatialStrategy(config);
      spatialStrategyCache.put(config.getRoundNumber(), strategy);
      //TODO remove previous round config?
      shapeConverter = makeShapeConverter(strategy, config, "doc.spatial.");
      System.out.println("Spatial Strategy: " + strategy);
    }
  }

  /**
   * Optionally converts points to circles, and optionally bbox'es result.
   */
  public static ShapeConverter makeShapeConverter(final SpatialStrategy spatialStrategy,
                                                  Config config, String configKeyPrefix) {
    //by default does no conversion
    final double radiusDegrees = config.get(configKeyPrefix+"radiusDegrees", 0.0);
    final double plusMinus = config.get(configKeyPrefix+"radiusDegreesRandPlusMinus", 0.0);
    final boolean bbox = config.get(configKeyPrefix + "bbox", false);

    return new ShapeConverter() {
      @Override
      public Shape convert(Shape shape) {
        if (shape instanceof Point && (radiusDegrees != 0.0 || plusMinus != 0.0)) {
          Point point = (Point)shape;
          double radius = radiusDegrees;
          if (plusMinus > 0.0) {
            Random random = new Random(point.hashCode());//use hashCode so it's reproducibly random
            radius += random.nextDouble() * 2 * plusMinus - plusMinus;
            radius = Math.abs(radius);//can happen if configured plusMinus > radiusDegrees
          }
          shape = spatialStrategy.getSpatialContext().makeCircle(point, radius);
        }
        if (bbox)
          shape = shape.getBoundingBox();
        return shape;
      }
    };
  }

  /** Converts one shape to another. Created by
   * {@link #makeShapeConverter(org.apache.lucene.spatial.SpatialStrategy, org.apache.lucene.benchmark.byTask.utils.Config, String)} */
  public interface ShapeConverter {
    Shape convert(Shape shape);
  }

  @Override
  public Document makeDocument() throws Exception {

    DocState docState = getDocState();

    Document doc = super.makeDocument();

    // Set SPATIAL_FIELD from body
    DocData docData = docState.docData;
    //   makeDocument() resets docState.getBody() so we can't look there; look in Document
    String shapeStr = doc.getField(DocMaker.BODY_FIELD).stringValue();
    Shape shape = makeShapeFromString(strategy, docData.getName(), shapeStr);
    if (shape != null) {
      shape = shapeConverter.convert(shape);
      //index
      for (Field f : strategy.createIndexableFields(shape)) {
        doc.add(f);
      }
    }

    return doc;
  }

  public static Shape makeShapeFromString(SpatialStrategy strategy, String name, String shapeStr) {
    if (shapeStr != null && shapeStr.length() > 0) {
      try {
        return strategy.getSpatialContext().readShapeFromWkt(shapeStr);
      } catch (Exception e) {//InvalidShapeException TODO
        System.err.println("Shape "+name+" wasn't parseable: "+e+"  (skipping it)");
        return null;
      }
    }
    return null;
  }

  @Override
  public Document makeDocument(int size) throws Exception {
    //TODO consider abusing the 'size' notion to number of shapes per document
    throw new UnsupportedOperationException();
  }
}