/*
 * 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.defaults.geo;

import java.awt.geom.AffineTransform;
import java.util.EnumSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;

import org.geotools.api.geometry.MismatchedDimensionException;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;
import org.geotools.api.referencing.datum.PixelInCell;
import org.geotools.api.referencing.operation.MathTransform;
import org.geotools.api.referencing.operation.TransformException;
import org.geotools.coverage.grid.GridEnvelope2D;
import org.geotools.coverage.grid.GridGeometry2D;
import org.geotools.referencing.CRS;
import org.geotools.referencing.operation.transform.ProjectiveTransform;
import org.geotools.util.factory.Hints;
import org.locationtech.jts.geom.Coordinate;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.geom.Point;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.geo.GeometryFamily;
import nz.org.riskscape.engine.geo.GeometryUtils;
import nz.org.riskscape.engine.grid.FeatureGrid;
import nz.org.riskscape.engine.grid.FeatureGridCell;

@Slf4j
@RequiredArgsConstructor
public class CutGeometryByGridOp {

  public enum Options {
    /**
     * Cutting a polygon against the grid cell polygon can result in 'glancing' intersections, e.g. if the two polygons
     * touch on one side, this can produce a point or a line string result. This options removes these mismatching
     * geometry types from the result returned by {@link #getCellIterator()}.
     * NB I can't think of cases where you would want this option turned off (maybe for debugging), but it's
     * there just in case.
     */
    MATCH_SOURCE_GEOMETRY
  }

  public static final EnumSet<Options> DEFAULT_OPTIONS = EnumSet.of(Options.MATCH_SOURCE_GEOMETRY);

  /**
   * Constructs a {@link GridGeometry2D} that can be used to segment geometry into pieces that are no
   * more then gridDistance in size in either direction. The grid (if it were extended if necessary) would
   * pass through alignGridTo with both the horizontal and vertical grid lines.
   *
   * Note: method allows package access for tests
   *
   * @param sridSet
   * @param geometry      geometry to create a grid for
   * @param gridDistance  the mesh size of the grid to create (in both directions)
   * @param alignGridTo   a coordinate to align the created grid to. If the created grid extended to cover
   *                      the alignGridTo point then both vertical and horizontal grid lines would pass through
   *                      this coordinate
   * @return
   */
  static GridGeometry2D constructGridGeometry(SRIDSet sridSet, Geometry geometry, double gridDistance,
      Coordinate alignGridTo) {
    CoordinateReferenceSystem crs = sridSet.get(geometry);

    Envelope featureEnvelope = geometry.getEnvelopeInternal();

    // We start by working out the delta (x and y) from where the feature's envelope is to one that is
    // aligned with alignGridTo
    double featureEnvelopeXDelta = (featureEnvelope.getMinX() - alignGridTo.getX()) % gridDistance;
    if (featureEnvelopeXDelta < 0) {
      // if the delta is less than zero then it is to the left of the alignGridTo. We need to mung it a little
      // to make it the delta from left of feature to grid boundary to the left of that
      featureEnvelopeXDelta = gridDistance + featureEnvelopeXDelta;
    }
    double featureEnvelopeYDelta = (featureEnvelope.getMinY() - alignGridTo.getY()) % gridDistance;
    if (featureEnvelopeYDelta < 0) {
      // if the delta is less than zero then it is below alignGridTo. We need to mung it a little
      // to make it the delta from bottom of feature to grid boundary below that.
      featureEnvelopeYDelta = gridDistance + featureEnvelopeYDelta;
    }
    double featureLowerLeftX = featureEnvelope.getMinX() - featureEnvelopeXDelta;
    double featureLowerLeftY = featureEnvelope.getMinY() - featureEnvelopeYDelta;

    // get the whole number of cells required to contain the feature
    double gridWidth = Math.ceil((featureEnvelope.getMaxX() - featureLowerLeftX) / gridDistance);
    double gridHeight = Math.ceil((featureEnvelope.getMaxY() - featureLowerLeftY) / gridDistance);

    // expand the envelope so that it will fit exactly the number of cells defined by the griddistance
    GridEnvelope2D gridRange = new GridEnvelope2D(0, 0, (int) gridWidth, (int) gridHeight);

    // now we reconstruct the featureEnvelope so it is origin aligned.
    featureEnvelope = new Envelope(featureLowerLeftX, featureLowerLeftX + gridWidth * gridDistance,
        featureLowerLeftY, featureLowerLeftY + gridHeight * gridDistance);

    // build an affine transform from the crs to the grid
    // NB this was cribbed from the GridGeometry2D code, but I'm not flipping the y axis, which makes debugging
    // slightly easier
    final double scaleX = featureEnvelope.getWidth()  / gridRange.getWidth();
    final double scaleY = featureEnvelope.getHeight() / gridRange.getHeight();
    final double transX = featureEnvelope.getMinX()   - gridRange.x * scaleX;
    final double transY = featureEnvelope.getMinY()   - gridRange.y * scaleY;


    // TODO it would be nice to get rid of the pixel centering stuff, but at the moment it's easier just to adjust for
    // it than to adjust and re-test the code in GriddedGeometry that relies on it
    final AffineTransform tr = new AffineTransform(scaleX, 0, 0, scaleY, transX, transY);
    tr.translate(0.5, 0.5); // Maps to pixel center
    MathTransform transform = ProjectiveTransform.create(tr);

    GridGeometry2D grid = new GridGeometry2D(gridRange, PixelInCell.CELL_CENTER, transform, crs, new Hints());
    assert(grid.getEnvelope2D().contains(featureEnvelope));
    return grid;
  }

  private final SRIDSet sridSet;

  private final EnumSet<Options> options;

  public CutGeometryByGridOp(SRIDSet sridSet) {
    this(sridSet, DEFAULT_OPTIONS);
  }

  /**
   * Cut the geometry up into a grid with mesh size of gridDistance. The grid is centred on the geometry
   * being cut which should reduce the number of pieces it is cut into.
   *
   * @param geometry      the geometry to cut up
   * @param gridDistance  the grid size
   * @return list of new geometries that have been cut to the grid
   */
  public List<Geometry> apply(Geometry geometry, double gridDistance) {
    Envelope geomEnv = geometry.getEnvelopeInternal();
    return apply(buildFeatureGrid(geometry, constructGridGeometry(sridSet, geometry, gridDistance,
        new Coordinate(geomEnv.getMinX(), geomEnv.getMinY()))));
  }

  /**
   * Cut the geometry up into a grid with mesh size of gridDistance.
   *
   * The grid is formed from the origin. The advantage of this is that many geometries can be cut to the
   * same grid.
   *
   * @param geometry      the geometry to cut up
   * @param gridDistance  the grid size
   * @param origin        the point to align the segmenting relative to
   * @return list of new geometries that have been cut to the grid
   */
  public List<Geometry> apply(Geometry geometry, double gridDistance, Point origin) {
    if (geometry.getSRID() != origin.getSRID()) {
      log.warn("Geometry (SRID {}) and origin (SRID {}) are in different CRSs", geometry.getSRID(), origin.getSRID());
    }

    return apply(buildFeatureGrid(geometry,
        constructGridGeometry(sridSet, geometry, gridDistance, origin.getCoordinate())));
  }

  private List<Geometry> apply(FeatureGrid featureGrid) {
    LinkedList<Geometry> collected = new LinkedList<>();

    Iterator<FeatureGridCell> cellIterator = featureGrid.cellIterator();

    GeometryFamily expectedFamily = GeometryFamily.from(featureGrid.getFeature());

    while (cellIterator.hasNext()) {
      FeatureGridCell cell = cellIterator.next();

      Geometry cut = cell.computeIntersection();

      if (options.contains(Options.MATCH_SOURCE_GEOMETRY)) {
        // NB: the intersection can turn parts of polygons into points or line-strings,
        // so check we're returning the same geometry type as what is being gridded
        cut = GeometryUtils.removeNonFamilyMembers(cut, expectedFamily);
      }

      if (!cut.isEmpty()) {
        // TODO unrolling multi-geometry like this is probably unnecessary, but we've done it like that historically
        GeometryUtils.processPerPart(cut, c -> collected.add(c));
      }
    }

    return collected;
  }

  private FeatureGrid buildFeatureGrid(Geometry polygon, GridGeometry2D gridGeometry) {
    try {
      return new FeatureGrid(polygon, CRS.getAxisOrder(sridSet.get(polygon)), gridGeometry);
    } catch (MismatchedDimensionException | TransformException e) {
      throw new RuntimeException("Unexpected error while building grid", e);
    }
  }


}
