/*
 * RiskScape™ Copyright New Zealand Institute for Earth Science Limited
 * (Earth Sciences New Zealand) is distributed for research purposes only
 * under the terms of AGPLv3.
 *
 * RiskScape™ Copyright 2026 New Zealand Institute for Earth Science
 * Limited (Earth Sciences New Zealand). All rights reserved. Source code
 * available under the AGPLv3.
 * 
 * This program is free software: you can redistribute it and/or modify it under
 *  the terms of the GNU Affero General Public License as published by the Free
 *  Software Foundation, either version 3 of the License, or (at your option) any
 *  later version.
 * 
 * This program is distributed for RESEARCH PURPOSES ONLY, in the hope that it will
 * be useful for research and education initiatives.
 * 
 * If you are not a researcher, or you are a researcher who wishes to use this
 * program on terms other than AGPLv3 (including those who wish to restrict the
 * distribution of any source code created using this program), please contact:
 * https://riskscape.org.nz
 * 
 * This program is distributed WITHOUT ANY WARRANTY; without even the implied
 * warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Affero General Public License for more details.  You should have received a copy
 * of the GNU Affero General Public License along with this program.  If not, see
 * <http://www.gnu.org/licenses/>.
 * 
 * By way of summary only, under the AGPLv3:
 *     • Permissions of this strongest copyleft license are conditioned
 *       on making available complete source code of licensed works and
 *       modifications, which include larger works using a licensed work,
 *       under the same license.
 *     • Copyright and license notices must be preserved.
 *     • Contributors provide an express grant of patent rights.
 *     • When a modified version is used to provide a service over a
 *       network, the complete source code of the modified version must be made
 *       available.
 */
package nz.org.riskscape.engine.geo;

import java.util.AbstractList;
import java.util.EnumMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;

import org.locationtech.jts.geom.Coordinate;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.geom.GeometryCollection;
import org.locationtech.jts.geom.GeometryFactory;
import org.locationtech.jts.geom.LinearRing;
import org.locationtech.jts.geom.Point;
import org.locationtech.jts.geom.Polygon;
import org.locationtech.jts.geom.impl.PackedCoordinateSequence;

import com.google.common.collect.Lists;

import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.util.SegmentedList;

/**
 * Alternative to CutGeometryByGridOp that builds a list of geometries that are intersected at a given grid size
 * by breaking a feature up recursively in order to 'throw away' as many useless large areas as possible in order to cut
 * down on the number of intersections required.  This capitalizes on the fact that most shapes are featureless except
 * for on their boundary.
 */
public class RecursiveQuadGridOp {

  /**
   * view-only list that combines the two lists from a {@link Result} without having to actually build a new list, so
   * that we can avoid storing the zillions of regular squares that are typically formed when gridding a polygon
   */
  static class CombinedList extends AbstractList<Geometry> {

    private final List<Geometry> cuts;
    private final BlockBasedSquareList squares;

    CombinedList(Result result) {
      this.cuts = result.cuts;
      this.squares = result.squares;
    }

    @Override
    public Geometry get(int index) {
      int cutsSize = cuts.size();
      if (index < cutsSize) {
        return cuts.get(index);
      } else {
        return squares.get(index - cutsSize);
      }
    }

    @Override
    public Iterator<Geometry> iterator() {
      Iterator<Geometry> cutsIter = cuts.iterator();
      Iterator<Polygon> squaresIter = squares.iterator();
      return new Iterator<Geometry>() {

        @Override
        public boolean hasNext() {
          return cutsIter.hasNext() || squaresIter.hasNext();
        }

        @Override
        public Geometry next() {
          if (cutsIter.hasNext()) {
            return cutsIter.next();
          }

          return squaresIter.next();
        };

      };
    }

    @Override
    public int size() {
      return cuts.size() + squares.size();
    }

  }
  @RequiredArgsConstructor
  private static final class QueueElement {
    /**
     * Envelope surrounding this geometry that's derived from the starting envelope
     */
    final Envelope envelope;
    /**
     * Geometry to break up
     */
    final Geometry geometry;

    /**
     * Recursive depth of this element - we track this integer value against a measured max depth so we know when to
     * stop recursing.  Originally I had used the width of the envelope and compared that to the target size, but
     * floating point inaccuracy meant that it didn't work reliably
     */
    final int depth;
  }

  /**
   * All the bits and pieces involved in the recursive quadding, plus ways of accessing it
   */
  public static class Result {

    /**
     * Size of the grid in map units, NB gridSize² gives the area of each grid cell.
     */
    public final double gridSize;

    /**
     * A geometry factory to use for creating new geometries from squares
     */
    public final GeometryFactory factory;

    /**
     * The irregular pieces of the original geometry that could not be cut into uniform squares
     */
    public final List<Geometry> cuts = SegmentedList.forClass(Geometry.class);

    /**
     * Uniform squares that comprise part of the original input geometry. Rather than cutting
     * grid-size pieces out of the original geometry (expensive!), we can construct the grid-size
     * pieces much more quickly from these larger squares. This works because the large squares
     * measure a power-of-two of the gridSize and are already aligned.
     */
    public final BlockBasedSquareList squares;

    /**
     * Maximum recursion depth - our envelope should match the grid size at this point.  NB we use a precomputed depth
     * to avoid having to do a fuzzy comparison of the envelope's size, which can be inaccurate with floating point
     * maths and storing numbers like 0.1 in binary
     */
    public final int maxDepth;

    private final Envelope startingEnvelope;

    public Result(Geometry geometry, double gridSize, Coordinate alignTo) {
      this.factory = geometry.getFactory();
      this.gridSize = gridSize;
      this.squares = new BlockBasedSquareList(factory, gridSize);
      this.startingEnvelope = geometry.getEnvelopeInternal();
      this.maxDepth = expandForFactor(startingEnvelope, gridSize, alignTo);
    }


    /**
     * @return a copy of the envelope that was determined to:
     * 1 - encompass the at least the entire geometry
     * 2 - be recursively divisible in to squares of size gridSize²
     * 3 - have its center 'snapped' to a grid defined by alignTo and gridSize
     */

    public Envelope getStartingEnvelope() {
      return new Envelope(startingEnvelope);
    }

    /**
     * The list of blocks we built when recursing, as polygons.  Can be useful for debugging/visualising how the
     * algorithm works
     */
    public List<? extends Geometry> getBlockList() {
      return Lists.transform(squares.getBlocks(), block -> new Polygon(envelopeToRing(factory, block), null, factory));
    }

    /**
     * A list of all the grid cell/feature intersections
     */
    public List<? extends Geometry> getCombinedResult() {
      return new CombinedList(this);
    }
  }

  /**
   * Adjust an envelope to use for the first intersection.  This will be a power-of-two multiple size of the given
   * required size, such that recursively partitioning the given envelope will eventually yield an envelope the size of
   * requiredSize
   * @return the max number of recursions
   */
  static int expandForFactor(Envelope featureEnvelope, double gridSize, Coordinate alignTo) {

    // base case - empty geometry gets no expansion
    if (featureEnvelope.isNull()) {
      return 0;
    }

    // make sure the target envelope is correctly aligned
    snapEnvelopeCentreTo(featureEnvelope, gridSize, alignTo);

    double target = Math.max(Math.max(featureEnvelope.getWidth(), featureEnvelope.getHeight()), gridSize);
    int exp = 0;
    double size = gridSize;
    // NB we need to be at least bigger than the target, otherwise we can end up in a situation (admittedly only likely
    // in tests) where the we are asked to split up a feature that is exactly half (or other power of 2) less than our
    // grid size, which means we start the recursive gridding with the first set of quads being half of the gridSize.
    // This matters if we end up creating a block (rather than a cut) from that size, which will end up too big.  Again,
    // this is only likely to happen in tests, but that's OK
    while (size <= target) {
      size = gridSize * (1 << ++exp);
    }

    Coordinate centre = featureEnvelope.centre();
    featureEnvelope.setToNull();
    featureEnvelope.expandToInclude(centre);
    featureEnvelope.expandBy(size / 2);

    return exp - 1;
  }

  /**
   * Modifies featureEnvelope to be snapped on to a grid defined by gridSize and alignTo, so that the envelope's centre
   * is now centered on that grid and the new envelope is at least as big as it was.  Note that this method doesn't
   * ensure that the envelope ends up in multiples of gridSize - that's the role of expandForFactor
   */
  static void snapEnvelopeCentreTo(Envelope featureEnvelope, double gridSize, Coordinate alignTo) {
    Coordinate centre = featureEnvelope.centre();
    Coordinate alignedCentre = new Coordinate(
        snapTo(centre.x, alignTo.x, gridSize),
        snapTo(centre.y, alignTo.y, gridSize)
    );

    // we need to increase the dimensions of the envelope so that the new envelope still covers the original one, but
    // its centre should now be alignedCentre
    double deltaX = Math.abs(centre.x - alignedCentre.x);
    double deltaY = Math.abs(centre.y - alignedCentre.y);
    double origWidth = featureEnvelope.getWidth();
    double origHeight = featureEnvelope.getHeight();

    featureEnvelope.setToNull();
    featureEnvelope.expandToInclude(alignedCentre);
    // we don't know which way the centre moved, so enlarge by the delta in each direction
    featureEnvelope.expandBy(origWidth / 2 + deltaX, origHeight / 2 + deltaY);

    assert featureEnvelope.centre().equals2D(alignedCentre);
  }

  static double snapTo(double snapping, double alignedTo, double snapSize) {
    // take the distance from the alignedTo origin and make it a multiple of snap size
    double totalDiff = snapping - alignedTo;
    double alignedTotalDiff = Math.round(totalDiff / snapSize) * snapSize;
    // the resulting delta is how much to move the original position
    double delta = totalDiff - alignedTotalDiff;
    return snapping - delta;
  }

  private boolean fastIntersectionEnabled = false;

  public List<? extends Geometry> apply(Geometry originalGeom, double gridSize) {
    return applyDetailed(originalGeom, gridSize).getCombinedResult();
  }

  public List<? extends Geometry> apply(Geometry originalGeom, double gridSize, Point alignTo) {
    return applyDetailed(originalGeom, gridSize, alignTo.getCoordinate()).getCombinedResult();
  }

  public List<? extends Geometry> apply(Geometry originalGeom, double gridSize, Coordinate alignTo) {
    return applyDetailed(originalGeom, gridSize, alignTo).getCombinedResult();
  }

  public Result applyDetailed(Geometry originalGeom, double gridSize) {
    // NB this is consistent with the CutGeometryByGridOp default
    Coordinate alignTo = Quadrant.BL.getCoordinate(originalGeom.getEnvelopeInternal());
    return applyDetailed(originalGeom, gridSize, alignTo);
  }


  /**
   * Compute the intersection of a grid of the given size with the given geometry, aligning the grid to the given
   * coordinate
   */
  public Result applyDetailed(Geometry originalGeom, double gridSize, Coordinate alignTo) {

    Result result = new Result(
      originalGeom,
      gridSize,
      alignTo
    );

    if (originalGeom.isEmpty()) {
      result.cuts.add(originalGeom);
      return result;
    }

    Envelope startingEnvelope = result.getStartingEnvelope();

    // NB we are doing BFS, but we could do a DFS - not sure which will end up performing better - not sure it'll make
    // any functional impact
    LinkedList<QueueElement> queue = new LinkedList<>();
    queue.add(new QueueElement(startingEnvelope, originalGeom, 0));

    while (!queue.isEmpty()) {
      QueueElement element = queue.removeFirst();

      // unwrap geometry collections to simplify logic
      if (element.geometry instanceof GeometryCollection) {
        GeometryCollection coll = (GeometryCollection) element.geometry;
        for (int i = 0; i < coll.getNumGeometries(); i++) {
          doIntersections(result, queue, new QueueElement(element.envelope, coll.getGeometryN(i), element.depth));
        }
      } else {
        doIntersections(result, queue, element);
      }
    }

    return result;
  }

  private void doIntersections(
      Result result,
      LinkedList<QueueElement> queue,
      QueueElement element
  ) {

    EnumMap<Quadrant, Envelope> envelopes = Quadrant.partition(element.envelope);

    // slow mode
    doIntersections(result, envelopes, queue, element);
  }

  /**
   * General purpose intersection method
   */
  private void doIntersections(Result result, EnumMap<Quadrant, Envelope> envelopes,
      LinkedList<QueueElement> queue, QueueElement element) {

    GeometryFamily targetFamily = GeometryFamily.from(element.geometry);
    for (Envelope envelope : envelopes.values()) {
      LinearRing envRing = envelopeToRing(element.geometry.getFactory(), envelope);
      Geometry intersection = element.geometry.intersection(new Polygon(envRing, null, envRing.getFactory()));

      // skip family member check, which avoids creating unnecessary empty geometries
      if (intersection.isEmpty()) {
        continue;
      }

      intersection = GeometryUtils.removeNonFamilyMembers(intersection, targetFamily);

      // skip this particular quadrant - it doesn't intersect the given geometry at all
      if (intersection.isEmpty()) {
        continue;
      }

      // NB there's a further optimization possible here - a polygon (or line) that only has points that lie on the
      // envelope doesn't need further breaking down - these can be stored a bit like the squares emitted specially.
      // This would improve performance very much on detailed grids with less detailed geometry
      LinearRing square = getPossibleSquare(intersection);

      // if the intersection completely overlaps this quadrant square, then we can stop
      // recursing. The BlockBasedSquareList can cut the geometry up further, as needed
      if (square != null) {
        if (isOverlapping(square, envRing)) {
          result.squares.addBlock(envelope);
          continue;
        }
      }

      if (element.depth < result.maxDepth) {
        // try recursing further to see if we can get smaller wholly-contained squares
        queue.add(new QueueElement(envelope, intersection, element.depth + 1));
      } else {
        // we're left with an non-square piece, e.g. the irregular coastal edge of a polygon
        result.cuts.add(intersection);
      }
    }
  }

  private boolean isOverlapping(LinearRing square, LinearRing envelopeAsRing) {
    // TODO we can do a much quicker check here - this is going to be unnecessarily complex, but it's probably not
    // a big deal in the grand scheme of things
    Geometry difference = envelopeAsRing.symDifference(square);
    return difference.isEmpty();
  }

  /**
   * @return a linear ring of size 5 (i.e. 4 unique points) from a given geometry,
   *         if it is a polygon with no holes
   */
  private LinearRing getPossibleSquare(Geometry geometry) {
    if (geometry instanceof Polygon) {
      Polygon polygon = (Polygon) geometry;

      if (polygon.getNumInteriorRing() == 0) {
        geometry = polygon.getExteriorRing();
      }
    }

    if (geometry instanceof LinearRing) {
      LinearRing ring = (LinearRing) geometry;
      if (ring.getCoordinateSequence().size() == 5) {
        return ring;
      }
    }

    return null;
  }

  /**
   * @return a LinearRing that represents the square defined by the given envelope
   */
  public static LinearRing envelopeToRing(GeometryFactory factory, Envelope envelope) {
    return new LinearRing(new PackedCoordinateSequence.Double(new Coordinate[] {
        Quadrant.TL.getCoordinate(envelope),
        Quadrant.TR.getCoordinate(envelope),
        Quadrant.BR.getCoordinate(envelope),
        Quadrant.BL.getCoordinate(envelope),
        Quadrant.TL.getCoordinate(envelope)
    }, 2, 0), factory);
  }
}
