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

import java.awt.image.Raster;

import javax.media.jai.PlanarImage;

import org.geotools.api.geometry.MismatchedDimensionException;
import org.geotools.api.geometry.Position;
import org.geotools.api.metadata.spatial.PixelOrientation;
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.GridCoverage2D;
import org.geotools.coverage.util.CoverageUtilities;
import org.geotools.geometry.Position2D;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.locationtech.jts.geom.Coordinate;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.GeometryFactory;

import it.geosolutions.jaiext.range.NoDataContainer;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.relation.BaseRelation;
import nz.org.riskscape.engine.relation.PeekingTupleIterator;
import nz.org.riskscape.engine.relation.TupleIterator;
import nz.org.riskscape.engine.types.CoverageType;
import nz.org.riskscape.engine.types.Referenced;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Types;

/**
 * Gives relation-style access to a grid by iteration through the raster pixel-by-pixel
 */
public class GridCoverageRelation extends BaseRelation {

  /**
   * The most minimal struct that a {@link GridCoverageRelation} will produce.  This is here solely as a best-efforts
   * link between this class and ToRelation so we can realize a dynamic expression with `to_relation`.
   */
  public static Struct createStructType(CoverageType coverageType) {
    return Struct.of("value", Types.FLOATING, "geom", Types.GEOMETRY);
  }

  public static GridCoverageRelation create(GridTypedCoverage typedCoverage, GridCoverage2D coverage) {

    CoordinateReferenceSystem crs = coverage.getCoordinateReferenceSystem2D();
    ReferencedEnvelope envelope2d = coverage.getEnvelope2D();
    Envelope envelope = new ReferencedEnvelope(envelope2d, crs);

    Struct struct = Struct.of("value", Types.FLOATING, "geom", Referenced.of(Types.POLYGON, crs, envelope));

    return new GridCoverageRelation(typedCoverage.getSridSet(), coverage, struct);
  }

  private final GridCoverage2D coverage;
  private final CoordinateReferenceSystem crs;
  private final SRIDSet sridSet;
  private NoDataContainer noData;

  protected GridCoverageRelation(SRIDSet sridSet, GridCoverage2D coverage, Struct type) {
    super(type);
    this.coverage = coverage;
    this.crs = coverage.getCoordinateReferenceSystem2D();
    this.sridSet = sridSet;
    this.noData = CoverageUtilities.getNoDataProperty(coverage);
  }

  protected GridCoverageRelation(SRIDSet sridSet, GridCoverage2D coverage, Fields fields) {
    super(fields);
    this.coverage = coverage;
    this.crs = coverage.getCoordinateReferenceSystem2D();
    this.sridSet = sridSet;
    this.noData = CoverageUtilities.getNoDataProperty(coverage);
  }

  @Override
  public String getSourceInformation() {
    return coverage.toString();
  }

  /**
   * Simple row based iteration through pixels.  TODO could be optimized to be tile-based and then do some work on
   * setting a decent tile size based on how image format is laid out on disk, e.g ascii grid should have height of 1
   * and geotiff should line up with the defined chunk size
   */
  private class IterImpl extends PeekingTupleIterator {

    private final Struct struct = getRawType();
    private final GeometryFactory gf = sridSet.getGeometryFactory(crs);
    private long index = 0;
    private int x = 0;
    private int y = 0;
    private final int pixelWidth = coverage.getRenderedImage().getWidth();
    private final int pixelHeight = coverage.getRenderedImage().getHeight();
    private final long totalPixels = pixelWidth * (long) pixelHeight;
    MathTransform gridToCrsLowerLeft = coverage.getGridGeometry().getGridToCRS(PixelOrientation.LOWER_LEFT);
    MathTransform gridToCrsLowerRight = coverage.getGridGeometry().getGridToCRS(PixelOrientation.LOWER_RIGHT);
    MathTransform gridToCrsUpperLeft = coverage.getGridGeometry().getGridToCRS(PixelOrientation.UPPER_LEFT);
    MathTransform gridToCrsUpperRight = coverage.getGridGeometry().getGridToCRS(PixelOrientation.UPPER_RIGHT);
    private double noDataValue = noData == null ? Double.NaN : noData.getAsSingleValue();
    // we know this to be true
    private final PlanarImage image = (PlanarImage) coverage.getRenderedImage();
    double[] resultPtr = new double[1];


    @Override
    public Tuple get() {
      while (index < totalPixels) {

        Position2D position = new Position2D(x, y);
        Position transformedLowerLeft;
        Position transformedLowerRight;
        Position transformedUpperLeft;
        Position transformedUpperRight;
        try {
          transformedLowerLeft = gridToCrsLowerLeft.transform(position, null);
          transformedLowerRight = gridToCrsLowerRight.transform(position, null);
          transformedUpperLeft = gridToCrsUpperLeft.transform(position, null);
          transformedUpperRight = gridToCrsUpperRight.transform(position, null);
        } catch (MismatchedDimensionException | TransformException e) {
          // TODO reuse same wrapping logic as in srid set?
          throw new RuntimeException(e);
        }
        int tileX = image.XToTileX(x);
        int tileY = image.YToTileY(y);

        Raster tile = image.getTile(tileX, tileY);
        if (tile == null) {
          // hmm, gap in the image data - seems like it's valid for there to be big holes in geotiffs
          increment();
          continue;
        }
        tile.getPixel(x, y, resultPtr);
        double result = resultPtr[0];

        //move indices for next call
        increment();

        if (result == noDataValue || Double.isNaN(noDataValue) && Double.isNaN(result)) {
          continue;
        }

        Tuple tuple = Tuple.ofValues(struct, resultPtr[0], gf.createPolygon(new Coordinate[] {
          new Coordinate(transformedLowerLeft.getOrdinate(0), transformedLowerLeft.getOrdinate(1)),
          new Coordinate(transformedLowerRight.getOrdinate(0), transformedLowerRight.getOrdinate(1)),
          new Coordinate(transformedUpperRight.getOrdinate(0), transformedUpperRight.getOrdinate(1)),
          new Coordinate(transformedUpperLeft.getOrdinate(0), transformedUpperLeft.getOrdinate(1)),
          new Coordinate(transformedLowerLeft.getOrdinate(0), transformedLowerLeft.getOrdinate(1))
        }));

        return tuple;
      }

      return null;
    }

    private void increment() {
      index++;
      x++;
      if (x >= pixelWidth) {
        x = 0;
        y++;
      }
    }

    @Override
    protected String getSource() {
      return coverage.toString();
    }
  }

  @Override
  protected TupleIterator rawIterator() {
    return new IterImpl();
  }

  @Override
  protected BaseRelation clone(Fields fields) {
    return new GridCoverageRelation(sridSet, coverage, fields);
  }

}
