/*
 * 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.spatial.prefix.tree;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.NoSuchElementException;
import org.apache.lucene.util.BytesRef;
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.locationtech.spatial4j.shape.impl.RectangleImpl;

/**
 * Uses a compact binary representation of 8 bytes to encode a spatial quad trie.
 *
 * <p>The binary representation is as follows:
 *
 * <pre>
 * CCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCCDDDDDL
 *
 * Where C = Cell bits (2 per quad)
 *       D = Depth bits (5 with max of 29 levels)
 *       L = isLeaf bit
 * </pre>
 *
 * It includes a built-in "pruneLeafyBranches" setting (true by default) similar to {@link
 * org.apache.lucene.spatial.prefix.RecursivePrefixTreeStrategy#setPruneLeafyBranches(boolean)}
 * although this one only prunes at the target detail level (where it has the most effect). Usually
 * you should disable RPT's prune, since it is very memory in-efficient.
 *
 * @lucene.experimental
 */
public class PackedQuadPrefixTree extends QuadPrefixTree {
  public static final int MAX_LEVELS_POSSIBLE = 29;
  protected static final byte[] QUAD = new byte[] {0x00, 0x01, 0x02, 0x03};

  protected boolean leafyPrune = true;

  /** Factory for creating {@link PackedQuadPrefixTree} instances with useful defaults. */
  public static class Factory extends QuadPrefixTree.Factory {
    @Override
    protected SpatialPrefixTree newSPT() {
      PackedQuadPrefixTree tree =
          new PackedQuadPrefixTree(ctx, maxLevels != null ? maxLevels : MAX_LEVELS_POSSIBLE);
      @SuppressWarnings("deprecation")
      Version lucene830 = Version.LUCENE_8_3_0;
      tree.robust = getVersion().onOrAfter(lucene830);
      return tree;
    }
  }

  public PackedQuadPrefixTree(SpatialContext ctx, int maxLevels) {
    super(ctx, maxLevels);
    if (maxLevels > MAX_LEVELS_POSSIBLE) {
      throw new IllegalArgumentException(
          "maxLevels of " + maxLevels + " exceeds limit of " + MAX_LEVELS_POSSIBLE);
    }
  }

  @Override
  public String toString() {
    return getClass().getSimpleName()
        + "(maxLevels:"
        + maxLevels
        + ",ctx:"
        + ctx
        + ",prune:"
        + leafyPrune
        + ")";
  }

  @Override
  public Cell getWorldCell() {
    return new PackedQuadCell(0x0L);
  }

  @Override
  public Cell getCell(Point p, int level) {
    if (!robust) { // old method
      List<Cell> cells = new ArrayList<>(1);
      buildNotRobustly(
          xmid, ymid, 0, cells, 0x0L, ctx.getShapeFactory().pointXY(p.getX(), p.getY()), level);
      if (!cells.isEmpty()) {
        return cells.get(0); // note cells could be longer if p on edge
      }
    }

    double currentXmid = xmid;
    double currentYmid = ymid;
    double xp = p.getX();
    double yp = p.getY();
    long term = 0L;
    int levelLimit = level > maxLevels ? maxLevels : level;
    SpatialRelation rel = SpatialRelation.CONTAINS;
    for (int lvl = 0; lvl < levelLimit; lvl++) {
      int quad = battenberg(currentXmid, currentYmid, xp, yp);
      double halfWidth = levelW[lvl + 1];
      double halfHeight = levelH[lvl + 1];
      switch (quad) {
        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:
      }
      // set bits for next level
      term |= (((long) (quad)) << (64 - ((lvl + 1) << 1)));
      // increment level
      term = ((term >>> 1) + 1) << 1;
    }
    return new PackedQuadCell(term, rel);
  }

  protected void buildNotRobustly(
      double x, double y, int level, List<Cell> matches, long term, Shape shape, int maxLevel) {
    double w = levelW[level] / 2;
    double h = levelH[level] / 2;

    // Z-Order
    // http://en.wikipedia.org/wiki/Z-order_%28curve%29
    checkBattenbergNotRobustly(QUAD[0], x - w, y + h, level, matches, term, shape, maxLevel);
    checkBattenbergNotRobustly(QUAD[1], x + w, y + h, level, matches, term, shape, maxLevel);
    checkBattenbergNotRobustly(QUAD[2], x - w, y - h, level, matches, term, shape, maxLevel);
    checkBattenbergNotRobustly(QUAD[3], x + w, y - h, level, matches, term, shape, maxLevel);
  }

  protected void checkBattenbergNotRobustly(
      byte quad,
      double cx,
      double cy,
      int level,
      List<Cell> matches,
      long term,
      Shape shape,
      int maxLevel) {
    // short-circuit if we find a match for the point (no need to continue recursion)
    if (shape instanceof Point && !matches.isEmpty()) return;
    double w = levelW[level] / 2;
    double h = levelH[level] / 2;

    SpatialRelation v = shape.relate(ctx.getShapeFactory().rect(cx - w, cx + w, cy - h, cy + h));

    if (SpatialRelation.DISJOINT == v) {
      return;
    }

    // set bits for next level
    term |= (((long) (quad)) << (64 - (++level << 1)));
    // increment level
    term = ((term >>> 1) + 1) << 1;

    if (SpatialRelation.CONTAINS == v || (level >= maxLevel)) {
      matches.add(new PackedQuadCell(term, v.transpose()));
    } else { // SpatialRelation.WITHIN, SpatialRelation.INTERSECTS
      buildNotRobustly(cx, cy, level, matches, term, shape, maxLevel);
    }
  }

  @Override
  public Cell readCell(BytesRef term, Cell scratch) {
    PackedQuadCell cell = (PackedQuadCell) scratch;
    if (cell == null) cell = (PackedQuadCell) getWorldCell();
    cell.readCell(term);
    return cell;
  }

  @Override
  public CellIterator getTreeCellIterator(Shape shape, int detailLevel) {
    if (detailLevel > maxLevels) {
      throw new IllegalArgumentException(
          "detailLevel:" + detailLevel + " exceed max: " + maxLevels);
    }
    return new PrefixTreeIterator(shape, (short) detailLevel);
  }

  public boolean isPruneLeafyBranches() {
    return leafyPrune;
  }

  /**
   * Like {@link
   * org.apache.lucene.spatial.prefix.RecursivePrefixTreeStrategy#setPruneLeafyBranches(boolean)}
   * but more memory efficient and only applies to the detailLevel, where it has the most effect.
   */
  public void setPruneLeafyBranches(boolean pruneLeafyBranches) {
    this.leafyPrune = pruneLeafyBranches;
  }

  /** See binary representation in the javadocs of {@link PackedQuadPrefixTree}. */
  protected class PackedQuadCell extends QuadCell {
    private long term;

    PackedQuadCell(long term) {
      super(null, 0, 0);
      this.term = term;
      this.b_off = 0;
      this.bytes = longToByteArray(this.term, new byte[8]);
      this.b_len = 8;
      readLeafAdjust();
    }

    PackedQuadCell(long term, SpatialRelation shapeRel) {
      this(term);
      this.shapeRel = shapeRel;
    }

    @Override
    protected void readCell(BytesRef bytes) {
      shapeRel = null;
      shape = null;
      this.bytes = bytes.bytes;
      this.b_off = bytes.offset;
      this.b_len = (short) bytes.length;
      this.term = longFromByteArray(this.bytes, bytes.offset);
      readLeafAdjust();
    }

    private final int getShiftForLevel(final int level) {
      return 64 - (level << 1);
    }

    public boolean isEnd(final int level, final int shift) {
      return (term != 0x0L && ((((0x1L << (level << 1)) - 1) - (term >>> shift)) == 0x0L));
    }

    /**
     * Get the next cell in the tree without using recursion. descend parameter requests traversal
     * to the child nodes, setting this to false will step to the next sibling. Note: This complies
     * with lexicographical ordering, once you've moved to the next sibling there is no
     * backtracking.
     */
    public PackedQuadCell nextCell(boolean descend) {
      final int level = getLevel();
      final int shift = getShiftForLevel(level);
      // base case: can't go further
      if ((!descend && isEnd(level, shift)) || isEnd(maxLevels, getShiftForLevel(maxLevels))) {
        return null;
      }
      long newTerm;
      final boolean isLeaf = (term & 0x1L) == 0x1L;
      // if descend requested && we're not at the maxLevel
      if ((descend && !isLeaf && (level != maxLevels)) || level == 0) {
        // simple case: increment level bits (next level)
        newTerm = ((term >>> 1) + 0x1L) << 1;
      } else { // we're not descending or we can't descend
        newTerm = term + (0x1L << shift);
        // we're at the last sibling...force descend
        if (((term >>> shift) & 0x3L) == 0x3L) {
          // adjust level for number popping up
          newTerm = ((newTerm >>> 1) - (Long.numberOfTrailingZeros(newTerm >>> shift) >>> 1)) << 1;
        }
      }
      return new PackedQuadCell(newTerm);
    }

    @Override
    protected void readLeafAdjust() {
      isLeaf = ((0x1L) & term) == 0x1L;
      if (getLevel() == getMaxLevels()) {
        isLeaf = true;
      }
    }

    @Override
    public BytesRef getTokenBytesWithLeaf(BytesRef result) {
      result = getTokenBytesNoLeaf(result);
      if (isLeaf()) {
        result.bytes[8 - 1] |= 0x1L; // set leaf
      }
      return result;
    }

    @Override
    public BytesRef getTokenBytesNoLeaf(BytesRef result) {
      if (result == null) {
        result = new BytesRef(8);
      } else if (result.bytes.length < 8) {
        result.bytes = new byte[8];
      }
      result.bytes = longToByteArray(this.term, result.bytes);
      result.offset = 0;
      result.length = 8;
      // no leaf
      result.bytes[8 - 1] &= ~1; // clear last bit (leaf bit)
      return result;
    }

    @Override
    public int compareToNoLeaf(Cell fromCell) {
      PackedQuadCell b = (PackedQuadCell) fromCell;
      // TODO clear last bit without the condition
      final long thisTerm = (((0x1L) & term) == 0x1L) ? term - 1 : term;
      final long fromTerm = (((0x1L) & b.term) == 0x1L) ? b.term - 1 : b.term;
      final int result = Long.compareUnsigned(thisTerm, fromTerm);
      assert Math.signum(result)
          == Math.signum(
              compare(
                  longToByteArray(thisTerm, new byte[8]),
                  0,
                  8,
                  longToByteArray(fromTerm, new byte[8]),
                  0,
                  8)); // TODO remove
      return result;
    }

    @Override
    public int getLevel() {
      int l = (int) ((term >>> 1) & 0x1FL);
      return l;
    }

    @Override
    protected Collection<Cell> getSubCells() {
      List<Cell> cells = new ArrayList<>(4);
      PackedQuadCell pqc =
          (new PackedQuadCell(((term & 0x1) == 0x1) ? this.term - 1 : this.term)).nextCell(true);
      cells.add(pqc);
      cells.add((pqc = pqc.nextCell(false)));
      cells.add((pqc = pqc.nextCell(false)));
      cells.add(pqc.nextCell(false));
      return cells;
    }

    @Override
    protected QuadCell getSubCell(Point p) {
      return (PackedQuadCell)
          PackedQuadPrefixTree.this.getCell(p, getLevel() + 1); // not performant!
    }

    @Override
    public boolean isPrefixOf(Cell c) {
      PackedQuadCell cell = (PackedQuadCell) c;
      return (this.term == 0x0L) || isInternalPrefix(cell);
    }

    protected boolean isInternalPrefix(PackedQuadCell c) {
      final int shift = 64 - (getLevel() << 1);
      return ((term >>> shift) - (c.term >>> shift)) == 0x0L;
    }

    protected long concat(byte postfix) {
      // extra leaf bit
      return this.term | (((long) (postfix)) << ((getMaxLevels() - getLevel() << 1) + 6));
    }

    /** Constructs a bounding box shape out of the encoded cell */
    @Override
    protected Rectangle makeShape() {
      double xmin = PackedQuadPrefixTree.this.xmin;
      double ymin = PackedQuadPrefixTree.this.ymin;
      int level = getLevel();

      byte b;
      for (short l = 0, i = 1; l < level; ++l, ++i) {
        b = (byte) ((term >>> (64 - (i << 1))) & 0x3L);

        switch (b) {
          case 0x00:
            ymin += levelH[l];
            break;
          case 0x01:
            xmin += levelW[l];
            ymin += levelH[l];
            break;
          case 0x02:
            break; // nothing really
          case 0x03:
            xmin += levelW[l];
            break;
          default:
            throw new RuntimeException("unexpected quadrant");
        }
      }

      double width, height;
      if (level > 0) {
        width = levelW[level - 1];
        height = levelH[level - 1];
      } else {
        width = gridW;
        height = gridH;
      }
      return new RectangleImpl(xmin, xmin + width, ymin, ymin + height, ctx);
    }

    private long fromBytes(byte b1, byte b2, byte b3, byte b4, byte b5, byte b6, byte b7, byte b8) {
      return ((long) b1 & 255L) << 56
          | ((long) b2 & 255L) << 48
          | ((long) b3 & 255L) << 40
          | ((long) b4 & 255L) << 32
          | ((long) b5 & 255L) << 24
          | ((long) b6 & 255L) << 16
          | ((long) b7 & 255L) << 8
          | (long) b8 & 255L;
    }

    private byte[] longToByteArray(long value, byte[] result) {
      for (int i = 7; i >= 0; --i) {
        result[i] = (byte) ((int) (value & 255L));
        value >>= 8;
      }
      return result;
    }

    private long longFromByteArray(byte[] bytes, int ofs) {
      assert bytes.length >= 8;
      return fromBytes(
          bytes[0 + ofs],
          bytes[1 + ofs],
          bytes[2 + ofs],
          bytes[3 + ofs],
          bytes[4 + ofs],
          bytes[5 + ofs],
          bytes[6 + ofs],
          bytes[7 + ofs]);
    }

    /** Used for debugging, this will print the bits of the cell */
    @Override
    public String toString() {
      StringBuilder s = new StringBuilder(64);
      final int numberOfLeadingZeros = Long.numberOfLeadingZeros(term);
      for (int i = 0; i < numberOfLeadingZeros; i++) {
        s.append('0');
      }
      if (term != 0) s.append(Long.toBinaryString(term));
      return s.toString();
    }
  } // PackedQuadCell

  /**
   * This is a streamlined version of TreeCellIterator, with built-in support to prune at
   * detailLevel (but not recursively upwards).
   */
  protected class PrefixTreeIterator extends CellIterator {
    private Shape shape;
    private PackedQuadCell thisCell;
    private PackedQuadCell nextCell;

    private short level;
    private final short detailLevel;
    private CellIterator pruneIter;

    PrefixTreeIterator(Shape shape, short detailLevel) {
      this.shape = shape;
      this.thisCell = ((PackedQuadCell) (getWorldCell())).nextCell(true);
      this.detailLevel = detailLevel;
      this.nextCell = null;
    }

    @Override
    public boolean hasNext() {
      if (nextCell != null) {
        return true;
      }
      SpatialRelation rel;
      // loop until we're at the end of the quad tree or we hit a relation
      while (thisCell != null) {
        rel = thisCell.getShape().relate(shape);
        if (rel == SpatialRelation.DISJOINT) {
          thisCell = thisCell.nextCell(false);
        } else { // within || intersects || contains
          thisCell.setShapeRel(rel);
          nextCell = thisCell;
          if (rel == SpatialRelation.WITHIN) {
            thisCell.setLeaf();
            thisCell = thisCell.nextCell(false);
          } else { // intersects || contains
            level = (short) (thisCell.getLevel());
            if (level == detailLevel || pruned(rel)) {
              thisCell.setLeaf();
              if (shape instanceof Point) {
                thisCell.setShapeRel(SpatialRelation.WITHIN);
                thisCell = null;
              } else {
                thisCell = thisCell.nextCell(false);
              }
              break;
            }
            thisCell = thisCell.nextCell(true);
          }
          break;
        }
      }
      return nextCell != null;
    }

    private boolean pruned(SpatialRelation rel) {
      int leaves;
      if (rel == SpatialRelation.INTERSECTS && leafyPrune && level == detailLevel - 1) {
        for (leaves = 0, pruneIter = thisCell.getNextLevelCells(shape);
            pruneIter.hasNext();
            pruneIter.next(), ++leaves)
          ;
        return leaves == 4;
      }
      return false;
    }

    @Override
    public Cell next() {
      if (nextCell == null) {
        if (!hasNext()) {
          throw new NoSuchElementException();
        }
      }
      // overriding since this implementation sets thisCell in hasNext
      Cell temp = nextCell;
      nextCell = null;
      return temp;
    }

    @Override
    public void remove() {
      // no-op
    }
  }
}
