package org.apache.lucene.spatial.prefix.tree;
import java.io.PrintStream;
import java.text.NumberFormat;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Locale;
import org.apache.lucene.util.Version;
import org.locationtech.spatial4j.context.SpatialContext;
import org.locationtech.spatial4j.shape.Point;
import org.locationtech.spatial4j.shape.Rectangle;
import org.locationtech.spatial4j.shape.Shape;
import org.locationtech.spatial4j.shape.SpatialRelation;
import org.apache.lucene.util.BytesRef;
public class QuadPrefixTree extends LegacyPrefixTree {
public static class Factory extends SpatialPrefixTreeFactory {
@Override
protected int getLevelForDistance(double degrees) {
return newSPT().getLevelForDistance(degrees);
}
@Override
protected SpatialPrefixTree newSPT() {
QuadPrefixTree tree = new QuadPrefixTree(ctx,
maxLevels != null ? maxLevels : MAX_LEVELS_POSSIBLE);
tree.robust = getVersion().onOrAfter(Version.LUCENE_8_3_0);
return tree;
}
}
public static final int MAX_LEVELS_POSSIBLE = 50;
public static final int DEFAULT_MAX_LEVELS = 12;
protected final double xmin;
protected final double xmax;
protected final double ymin;
protected final double ymax;
protected final double xmid;
protected final double ymid;
protected final double gridW;
public final double gridH;
final double[] levelW;
final double[] levelH;
protected boolean robust = true;
public QuadPrefixTree(
SpatialContext ctx, Rectangle bounds, int maxLevels) {
super(ctx, maxLevels);
this.xmin = bounds.getMinX();
this.xmax = bounds.getMaxX();
this.ymin = bounds.getMinY();
this.ymax = bounds.getMaxY();
levelW = new double[maxLevels + 1];
levelH = new double[maxLevels + 1];
gridW = xmax - xmin;
gridH = ymax - ymin;
this.xmid = xmin + gridW/2.0;
this.ymid = ymin + gridH/2.0;
levelW[0] = gridW/2.0;
levelH[0] = gridH/2.0;
for (int i = 1; i < levelW.length; i++) {
levelW[i] = levelW[i - 1] / 2.0;
levelH[i] = levelH[i - 1] / 2.0;
}
}
public QuadPrefixTree(SpatialContext ctx) {
this(ctx, DEFAULT_MAX_LEVELS);
}
public QuadPrefixTree(
SpatialContext ctx, int maxLevels) {
this(ctx, ctx.getWorldBounds(), maxLevels);
}
@Override
public Cell getWorldCell() {
return new QuadCell(BytesRef.EMPTY_BYTES, 0, 0);
}
public void printInfo(PrintStream out) {
NumberFormat nf = NumberFormat.getNumberInstance(Locale.ROOT);
nf.setMaximumFractionDigits(5);
nf.setMinimumFractionDigits(5);
nf.setMinimumIntegerDigits(3);
for (int i = 0; i < maxLevels; i++) {
out.println(i + "]\t" + nf.format(levelW[i]) + "\t" + nf.format(levelH[i]));
}
}
@Override
public int getLevelForDistance(double dist) {
if (dist == 0)
return maxLevels;
for (int i = 0; i < maxLevels-1; i++) {
if(dist > levelW[i] && dist > levelH[i]) {
return i+1;
}
}
return maxLevels;
}
@Override
public Cell getCell(Point p, int level) {
if (!robust) {
List<Cell> cells = new ArrayList<>(1);
buildNotRobustly(xmid, ymid, 0, cells, new BytesRef(maxLevels+1), ctx.makePoint(p.getX(),p.getY()), level);
if (!cells.isEmpty()) {
return cells.get(0);
}
}
double currentXmid = xmid;
double currentYmid = ymid;
double xp = p.getX();
double yp = p.getY();
BytesRef str = new BytesRef(maxLevels+1);
int levelLimit = level > maxLevels ? maxLevels : level;
SpatialRelation rel = SpatialRelation.CONTAINS;
for (int lvl = 0; lvl < levelLimit; lvl++){
int c = battenberg(currentXmid, currentYmid, xp, yp);
double halfWidth = levelW[lvl + 1];
double halfHeight = levelH[lvl + 1];
switch(c){
case 0:
currentXmid -= halfWidth;
currentYmid += halfHeight;
break;
case 1:
currentXmid += halfWidth;
currentYmid += halfHeight;
break;
case 2:
currentXmid -= halfWidth;
currentYmid -= halfHeight;
break;
case 3:
currentXmid += halfWidth;
currentYmid -= halfHeight;
break;
default:
}
str.bytes[str.length++] = (byte)('A' + c);
}
return new QuadCell(str, rel);
}
private void buildNotRobustly(
double x,
double y,
int level,
List<Cell> matches,
BytesRef str,
Shape shape,
int maxLevel) {
assert str.length == level;
double w = levelW[level] / 2;
double h = levelH[level] / 2;
checkBattenbergNotRobustly('A', x - w, y + h, level, matches, str, shape, maxLevel);
checkBattenbergNotRobustly('B', x + w, y + h, level, matches, str, shape, maxLevel);
checkBattenbergNotRobustly('C', x - w, y - h, level, matches, str, shape, maxLevel);
checkBattenbergNotRobustly('D', x + w, y - h, level, matches, str, shape, maxLevel);
}
protected void checkBattenbergNotRobustly(
char c,
double cx,
double cy,
int level,
List<Cell> matches,
BytesRef str,
Shape shape,
int maxLevel) {
assert str.length == level;
assert str.offset == 0;
double w = levelW[level] / 2;
double h = levelH[level] / 2;
int strlen = str.length;
Rectangle rectangle = ctx.makeRectangle(cx - w, cx + w, cy - h, cy + h);
SpatialRelation v = shape.relate(rectangle);
if (SpatialRelation.CONTAINS == v) {
str.bytes[str.length++] = (byte)c;
matches.add(new QuadCell(BytesRef.deepCopyOf(str), v.transpose()));
} else if (SpatialRelation.DISJOINT == v) {
} else {
str.bytes[str.length++] = (byte)c;
int nextLevel = level+1;
if (nextLevel >= maxLevel) {
matches.add(new QuadCell(BytesRef.deepCopyOf(str), v.transpose()));
} else {
buildNotRobustly(cx, cy, nextLevel, matches, str, shape, maxLevel);
}
}
str.length = strlen;
}
protected int battenberg(double xmid, double ymid, double xp, double yp){
if (ymid <= yp) {
if (xmid >= xp) {
return 0;
}
return 1;
} else {
if (xmid >= xp) {
return 2;
}
return 3;
}
}
protected class QuadCell extends LegacyCell {
QuadCell(byte[] bytes, int off, int len) {
super(bytes, off, len);
}
QuadCell(BytesRef str, SpatialRelation shapeRel) {
this(str.bytes, str.offset, str.length);
this.shapeRel = shapeRel;
}
@Override
protected QuadPrefixTree getGrid() { return QuadPrefixTree.this; }
@Override
protected int getMaxLevels() { return maxLevels; }
@Override
protected Collection<Cell> getSubCells() {
BytesRef source = getTokenBytesNoLeaf(null);
List<Cell> cells = new ArrayList<>(4);
cells.add(new QuadCell(concat(source, (byte)'A'), null));
cells.add(new QuadCell(concat(source, (byte)'B'), null));
cells.add(new QuadCell(concat(source, (byte)'C'), null));
cells.add(new QuadCell(concat(source, (byte)'D'), null));
return cells;
}
protected BytesRef concat(BytesRef source, byte b) {
final byte[] buffer = new byte[source.length + 2];
System.arraycopy(source.bytes, source.offset, buffer, 0, source.length);
BytesRef target = new BytesRef(buffer);
target.length = source.length;
target.bytes[target.length++] = b;
return target;
}
@Override
public int getSubCellsSize() {
return 4;
}
@Override
protected QuadCell getSubCell(Point p) {
return (QuadCell) QuadPrefixTree.this.getCell(p, getLevel() + 1);
}
@Override
public Shape getShape() {
if (shape == null)
shape = makeShape();
return shape;
}
protected Rectangle makeShape() {
BytesRef token = getTokenBytesNoLeaf(null);
double xmin = QuadPrefixTree.this.xmin;
double ymin = QuadPrefixTree.this.ymin;
for (int i = 0; i < token.length; i++) {
byte c = token.bytes[token.offset + i];
switch (c) {
case 'A':
ymin += levelH[i];
break;
case 'B':
xmin += levelW[i];
ymin += levelH[i];
break;
case 'C':
break;
case 'D':
xmin += levelW[i];
break;
default:
throw new RuntimeException("unexpected char: " + c);
}
}
int len = token.length;
double width, height;
if (len > 0) {
width = levelW[len-1];
height = levelH[len-1];
} else {
width = gridW;
height = gridH;
}
return ctx.makeRectangle(xmin, xmin + width, ymin, ymin + height);
}
}
}