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

import java.util.Collections;
import java.util.Iterator;

import org.geotools.api.geometry.MismatchedDimensionException;
import org.geotools.api.geometry.Position;
import org.geotools.api.referencing.crs.CoordinateReferenceSystem;
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.geometry.Position2D;
import org.geotools.referencing.CRS.AxisOrder;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.geom.GeometryFactory;

import it.geosolutions.imageio.plugins.arcgrid.AsciiGridsImageMetadata.RasterSpaceType;
import lombok.Getter;

/**
 * Surrounds a geometry (called a feature in this context to help distinguish it from other uses of the word geometry)
 * with a grid to allow it to be cut or sampled in to cells as defined by the grid.
 */
public class FeatureGrid {

  private static Envelope buildGridEnvelope(Geometry feature, GridGeometry2D gridGeometry)
      throws MismatchedDimensionException, TransformException {
    MathTransform crsToGrid = gridGeometry.getCRSToGrid2D();

    if (! feature.getEnvelopeInternal().intersects(gridGeometry.getEnvelope2D())) {
      // no intersection here, return an empty envelope and move on
      return new Envelope();
    }
    // get the envelope surrounding the parts of the feature that is in the grid
    // clipping to the grid envelope is important as `crsToGrid.transform` may give bad answers if
    // you ask it to transform positions from outside of the grid
    Envelope featureInGridEnvelope = feature.getEnvelopeInternal().intersection(gridGeometry.getEnvelope2D());

    Position featureGridMin =
        crsToGrid.transform(new Position2D(featureInGridEnvelope.getMinX(), featureInGridEnvelope.getMinY()), null);

    Position featureGridMax =
        crsToGrid.transform(new Position2D(featureInGridEnvelope.getMaxX(), featureInGridEnvelope.getMaxY()), null);

    GridEnvelope2D gridRange2D = gridGeometry.getGridRange2D();

    // User awt .getMinY() which returns maximal values which expands the grid size for us.
    // refer to caution at: https://docs.geotools.org/latest/javadocs/org/geotools/coverage/grid/GridEnvelope2D.html
    Envelope coverageEnvelope = new Envelope(
        gridRange2D.getMinX(),
        gridRange2D.getMaxX(),
        gridRange2D.getMinY(),
        gridRange2D.getMaxY()
    );

    // Note that this method of building the grid envelope returns the smallest envelope possible.
    // this is only really significant for line (or point) features that run along grid lines.
    // when this occurs the envelope will be the cell above or to the right of the grid line.
    double minX = Math.floor(getMin(featureGridMin, featureGridMax, 0) + HALF_A_PIXEL);
    double maxX = Math.ceil(getMax(featureGridMin, featureGridMax, 0) + HALF_A_PIXEL);
    double minY = Math.floor(getMin(featureGridMin, featureGridMax, 1) + HALF_A_PIXEL);
    double maxY = Math.ceil(getMax(featureGridMin, featureGridMax, 1) + HALF_A_PIXEL);
    // for point/lines that are on the grid lines we can get the case that our min/max values
    // end up the same (but we need them to create an envelope that covers some number of grid cells
    // so if that happens we increase the max by one (this only actually happens in the rare case that
    // the featureGrid min/max are the same and x or y is 0.5 because that plus half a pixels is a whole
    // number so floor/ceil are the same whole number)
    if (minX == maxX) {
      maxX += 1D;
    }
    if (minY == maxY) {
      maxY += 1D;
    }

    return coverageEnvelope.intersection(new Envelope(minX, maxX, minY, maxY));
  }

  private static double getMax(Position featureGridMin, Position featureGridMax, int i) {
    return Math.max(featureGridMax.getOrdinate(i), featureGridMin.getOrdinate(i));
  }

  private static double getMin(Position featureGridMin, Position featureGridMax, int i) {
    return Math.min(featureGridMax.getOrdinate(i), featureGridMin.getOrdinate(i));
  }

  /**
   * GridGeometry2D typically models the grid with {@link RasterSpaceType#PixelIsArea}, so we must shunt our
   * coordinates by half at various times during gridding
   */
  public static final double HALF_A_PIXEL = 0.5;

  /**
   * The geometry that is being surrounded by the grid
   */
  @Getter
  private final Geometry feature;

  /**
   * The {@link AxisOrder} of feature.
   */
  private final AxisOrder axisOrder;

  /**
   * The geometry of the overall grid, e.g. this is the grid that we are applying over the top of our feature.
   */
  @Getter
  private final GridGeometry2D gridGeometry;

  /**
   * The envelope of the feature in grid space, snapped to the grid, so that is the smallest possible size that
   * surrounds the feature.
   */
  @Getter
  private final Envelope envelope;

  /**
   * A maths transformation from grid coordinates to world crs space
   */
  @Getter
  private final MathTransform gridToWorld;

  /**
   * Can be used for creating new geometries that are in the same crs as `feature`
   */
  private final GeometryFactory gf;

  /**
   * The number of columns in the grid, i.e. the width, formed by the intersection of the feature and the grid geometry.
   */
  @Getter
  private final int gridColumns;

  /**
   * The number of rows in the grid, i.e. the height, formed by the intersection of the feature and the grid geometry.
   */
  @Getter
  private final int gridRows;

  /**
   * The y position in the original grid where this feature grid starts
   */
  @Getter
  private final int gridOffsetY;

  /**
   * The x position in the original grid where this feature grid starts
   */
  @Getter
  private final int gridOffsetX;

  /**
   * The number of crs units between grid cells in the horizontal direction, e.g. you can add `cellWidth` to
   * `featureTopLeft` to move across one grid cell column.
   */
  @Getter
  private final double cellWidth;

  /**
   * The number of crs units between grid cells in the vertical direction, e.g. you can add `cellHeight` to
   * `featureTopLeft` to move down one grid cell row.
   */
  @Getter
  private final double cellHeight;

  /**
   * The top left position of the feature within the grid in crs space.  The is 'snapped' to the grid geometry and is
   * the maximum position it can be while still encompassing the entire feature.
   */
  @Getter
  private final Position2D featureTopLeft;

  /**
   * Construct a new {@link FeatureGrid} that applies a {@link GridGeometry2D} to a {@link Geometry} object.  No
   * re-projection is done by this class, it is up to the caller to ensure that the {@link CoordinateReferenceSystem} of
   * the gridGeometry's `world` matches the {@link CoordinateReferenceSystem} of the feature.
   */
  public FeatureGrid(Geometry feature, AxisOrder axisOrder, GridGeometry2D gridGeometry)
      throws MismatchedDimensionException, TransformException {

    this.feature = feature;
    this.axisOrder = axisOrder;
    this.gridGeometry = gridGeometry;
    this.gridToWorld = gridGeometry.getGridToCRS2D();
    this.gf = feature.getFactory();

    this.envelope = buildGridEnvelope(feature, gridGeometry);

    gridColumns = (int) envelope.getWidth();
    gridRows = (int) envelope.getHeight();
    gridOffsetY = (int) envelope.getMinY();
    gridOffsetX = (int) envelope.getMinX();

    cellWidth = computeCellDeltaInWorldSpace(0);
    cellHeight = computeCellDeltaInWorldSpace(1);
    featureTopLeft = computeTopLeftInWorldSpace();
  }

  public GeometryFactory getGeometryFactory() {
    return gf;
  }

  /**
   * @return an Envelope based on convering the grid's dimensions to the world's - note that this will not add the extra
   * pixel on to width and height that the geometry does - this method is really here for tests and sanity checking
   */
  public Envelope getFeatureGridEnvelopeInWorldCrs() {
    try {
      Position minCrs = gridToWorld.transform(new Position2D(
          envelope.getMinX() - HALF_A_PIXEL, envelope.getMinY() - HALF_A_PIXEL), null);

      Position maxCrs = gridToWorld.transform(new Position2D(
          envelope.getMaxX() - HALF_A_PIXEL, envelope.getMaxY() - HALF_A_PIXEL), null);

      return new Envelope(minCrs.getOrdinate(0), maxCrs.getOrdinate(0), minCrs.getOrdinate(1), maxCrs.getOrdinate(1));
    } catch (MismatchedDimensionException | TransformException e) {
      throw new RuntimeException("this shouldn't happen once construction validates the geometries", e);
    }
  }

  /**
   * Iterate over the points of the grid formed by intersection of the feature's bounds and the grid.
   */
  public Iterator<FeatureGridCell> cellIterator() {
    if (gridColumns == 0 || gridRows == 0) {
      // there are no cells to iterate of columns or rows are zero
      return Collections.emptyIterator();
    }
    return new FeatureGridIterator(this);
  }

  private Position2D computeTopLeftInWorldSpace() throws TransformException {
    MathTransform gridToCRS = gridGeometry.getGridToCRS2D();

    Position gridPoint1 = new Position2D(gridOffsetX, gridOffsetY);
    Position2D crsPoint1 = new Position2D();

    gridToCRS.transform(gridPoint1, crsPoint1);

    return crsPoint1;
  }

  private double computeCellDeltaInWorldSpace(int ordinate) throws TransformException {
    MathTransform gridToCRS = gridGeometry.getGridToCRS2D();

    Position gridPoint1 = new Position2D(gridOffsetX, gridOffsetY);
    Position gridPoint2 = new Position2D(gridOffsetX, gridOffsetY);

    // shift either across (0) or down (1)
    gridPoint2.setOrdinate(ordinate, gridPoint2.getOrdinate(ordinate) + 1);

    // transform both in to world space
    Position crsPoint1 = gridToCRS.transform(gridPoint1, null);
    Position crsPoint2 = gridToCRS.transform(gridPoint2, null);

    // the grid is always in a xy coordinate order. but the features CRS may to the other way around.
    int worldOrdinate = axisOrder == AxisOrder.EAST_NORTH ? ordinate : 1 - ordinate;
    // measure the difference
    return crsPoint2.getOrdinate(worldOrdinate) - crsPoint1.getOrdinate(worldOrdinate);
  }

  /**
   * @return the cell dimension for the x axis in the World CRS. This will be cellWidth or cellHeight
   * depending on the CRS axis order
   */
  protected double getCellWorldX() {
    return axisOrder == AxisOrder.EAST_NORTH ? cellWidth : cellHeight;
  }

  /**
   * @return the cell dimension for the y axis in the World CRS. This will be cellWidth or cellHeight
   * depending on the CRS axis order
   */
  protected double getCellWorldY() {
    return axisOrder == AxisOrder.EAST_NORTH ? cellHeight : cellWidth;
  }
}
