/*
 * 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 static nz.org.riskscape.engine.geo.GeometryUtils.*;

import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedList;
import java.util.List;
import java.util.Optional;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Stream;

import org.geotools.api.referencing.crs.CoordinateReferenceSystem;
import org.geotools.geometry.jts.ReferencedEnvelope;
import org.geotools.referencing.factory.epsg.CartesianAuthorityFactory;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.geom.LinearRing;
import org.locationtech.jts.geom.Point;
import org.locationtech.jts.geom.Polygon;
import org.locationtech.jts.index.strtree.STRtree;

import com.google.common.collect.Lists;

import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import nz.org.riskscape.engine.RiskscapeException;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.data.coverage.SpatialRelationTypedCoverage;
import nz.org.riskscape.engine.expr.StructMemberAccessExpression;
import nz.org.riskscape.engine.problem.GeneralProblems;
import nz.org.riskscape.engine.problem.ProblemFactory;
import nz.org.riskscape.engine.query.TupleUtils;
import nz.org.riskscape.engine.relation.Relation;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.Referenced;
import nz.org.riskscape.engine.types.Struct.StructMember;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.util.Pair;
import nz.org.riskscape.problem.Problem;
import nz.org.riskscape.problem.ProblemException;
import nz.org.riskscape.problem.Problems;
import nz.org.riskscape.problem.StandardCodes;


/**
 * Creates an index of tuples using an underlying {@link STRtree}, implementing a couple of standard routines for
 * querying the index for tuples, dealing with reprojection and other housekeeping concerns.
 *
 * Note that testing is covered off for this by {@link SpatialRelationTypedCoverage}.
 * (in defaults) but we could merge the tests and avoid retesting
 */
@RequiredArgsConstructor
@Slf4j
public class IntersectionIndex {

  public interface LocalProblems extends ProblemFactory {
    Problem intersectionCutNotSupportedHint(GeometryFamily geomClazz);
  }

  public static final LocalProblems PROBLEMS = Problems.get(LocalProblems.class);

  @Data
  public static class Options {

    /**
     * Whether to optimize this index by cutting up complex geometries for faster point lookups.  This optimization cuts
     * up the indexed geometry recursively in to quads, until all the constituent parts are no longer "too complex"
     * according to the indexes given options.  Note that this option can have a slight detrimental affect on the
     * performance of sampling polygonal features against the index.
     *
     * This improves performance by excluding large areas of these indexed geometries before they go in to the index, so
     * that each feature's indexed "hitbox" is smaller, meaning less false positives in the index.  This reduces the
     * number of times the index returns a hit, only for the later, more computationally expensive, intersection test to
     * fail.
     *
     * As a further optimization, point intersections can be done against the cut geometry, instead of the larger source
     * geometry, which runs quicker against smaller (in terms of points) features.
     *
     * This feature will increase the amount of memory used by the index.  Tuning the other options may help
     * reach a compromise between memory usage and performance - the defaults are meant to be a good compromise.
     *
     * If this option has been set to empty then the index will choose based on the first encountered geometry type
     * (polygons will be cut, lines and points will not).
     */
    private Optional<Boolean> cutBeforeAdding = Optional.of(false);

    /**
     * Any geometries with more points than this is too complex
     */
    private long cutPoints = 1000;

    /**
     * Any geometries with more points than this and a larger cutRatio is too complex
     */
    private long cutRatioPoints = 250;

    /**
     * Any geometries with more than cutMinPoints is too complex if the ratio of their area to their bounding box is
     * less than `cutRatio`.  This is a measure of the 'efficiency' of the bounding box as an index key;  The
     * higher the ratio, the more chance that a bounding box hit will also intersect with the indexed geometry.
     */
    private double cutRatio = 0.75;

    /**
     * An optional width/height to target for complexity.  When this option is set, this takes precedence over all other
     * options.  It defines the size of the geometry as being the (width + height) / 2.  Note that it won't do any
     * reprojection if your geometries are not metric.
     */
    private long cutSizeMapUnits = -1L;

    /**
     * Heuristic for deciding if geometry is 'too complex' to put in the index.  We use this to decide whether to cut it
     * up before adding it (as well as computing the true intersection for points later)
     */
    public boolean isTooComplex(Geometry geom, Envelope geomEnvelope) {
      int numPoints = geom.getNumPoints();

      // cut it if on average the bounding box is greater than a cut size
      if (cutSizeMapUnits != -1L) {

        if (geom instanceof Polygon poly) {
          if (poly.getNumGeometries() == 1) {
            // this is probably a square - even if it's "too big" there's no point cutting it further
            if (poly.getNumPoints() <= 5) {
              return false;
            }
          }
        }

        double averageSize = (geomEnvelope.getWidth() + geomEnvelope.getHeight()) / 2;
        if (averageSize > cutSizeMapUnits) {
          return true;
        }
      }

      // complex, too big, regardless of ratio
      if (numPoints > cutPoints) {
        return true;
      }

      // it's in the middle, see if the ratio of the area to the envelope is too small
      return numPoints > cutRatioPoints && geom.getArea() / geomEnvelope.getArea() < cutRatio;
    }

  }

  public static Options defaultOptions() {
    return new Options();
  }
  /**
   * Builds and populates an {@link IntersectionIndex} with the content of a {@link Relation}.
   *
   * If the {@link CoordinateReferenceSystem} of the rhsRelation and the lhsGeomType are different then
   * the index will be populated with the geometry member re-projected to the lhsGeomType's
   * projection if this is safe to do.
   *
   * In cases where there should be more lhs values than rhs then this should improve performance as there
   * would be less re-projection overall.
   *
   * Re-projecting the rhs is only safe to do if we know that all of it's features will fit within the
   * bounds of the lhs projection.
   *
   * @param rhsRelation relation to populate the index from
   * @param lhsGeomType the type of the geometry member that will be used when querying the index
   * @param sridSet
   * @return index that is populated from rhsRelation
   * @throws ProblemException if rhsRelation has no geometry member
   */
  public static IntersectionIndex populateFromRelation(
      Relation rhsRelation,
      Type lhsGeomType,
      SRIDSet sridSet,
      Options options
  ) throws ProblemException {
    StructMember geomMember;
    try {
      geomMember = TupleUtils.findGeometryMember(rhsRelation.getType(), TupleUtils.FindOption.ANY_REQUIRED);
    } catch (RiskscapeException e) {
      throw new ProblemException(Problem.error(StandardCodes.GEOMETRY_REQUIRED, rhsRelation.getType()));
    }

    IntersectionIndex index = new IntersectionIndex(geomMember, sridSet, options);

    Function<Tuple,Tuple> tupleMapper;
    Optional<Referenced> lhsReferencedGeomType = Nullable.strip(lhsGeomType).find(Referenced.class);
    if (lhsReferencedGeomType.isPresent() && canReprojectSafely(geomMember.getType(), lhsGeomType)) {
      int lhsSrid = sridSet.get(lhsReferencedGeomType.get().getCrs());
      tupleMapper = t -> {
        Geometry geom = t.fetch(geomMember);
        t.set(geomMember, sridSet.reproject(geom, lhsSrid));
        return t;
      };
    } else {
      tupleMapper = Function.identity();
    }

    rhsRelation.iterator().forEachRemaining(tuple -> {
      index.insert(tupleMapper.apply(tuple));
    });
    index.build();

    return index;
  }

  private final STRtree tree = new STRtree();

  /**
   * Total envelope of all indexed features
   */
  private final Envelope envelope = new Envelope();

  /**
   * The geometry member to use for accessing geometries from indexed tuples
   */
  private final StructMemberAccessExpression geometryMemberAccessor;

  /**
   * An {@link SRIDSet} for handling reprojection
   */
  private final SRIDSet sridSet;

  private final Options options;

  private final OverlayOperations overlayOperations = OverlayOperations.get();

  // The crs of indexed tuples - default to a wildcard crs - this should be pretty harmless in the case of an empty
  // index - we should be able to replace this once #285 has been sorted
  private CoordinateReferenceSystem crs = CartesianAuthorityFactory.GENERIC_2D;
  private int srid = -1;
  private GeometryFamily indexedFamily = GeometryFamily.PUNTAL;

  public static IntersectionIndex withDefaultOptions(StructMember expression, SRIDSet sridSet) {
    return new IntersectionIndex(expression, sridSet, new Options());
  }
  public static IntersectionIndex withDefaultOptions(StructMemberAccessExpression expression, SRIDSet sridSet) {
    return new IntersectionIndex(expression, sridSet, new Options());
  }

  public IntersectionIndex(StructMember geometryMember, SRIDSet sridSet, Options options) {
    this.geometryMemberAccessor = new StructMemberAccessExpression(false, geometryMember);
    this.sridSet = sridSet;
    this.options = options;
  }

  /**
   * Builds the backing {@link STRtree}. Once built then no further tuples should be inserted.
   *
   * The index should be built as soon as all tuples have been inserted from a single threaded context.
   * This is to prevent potential NPE from the backing STRtree (refer to GL#1547).
   */
  public void build() {
    tree.build();
  }

  public void insert(Tuple tuple) {
    Geometry geom = geometryMemberAccessor.evaluate(tuple, Geometry.class);
    if (geom == null) {
      // null geometries are not indexed
      return;
    }

    if (srid == -1) {
      // this is the first geometry to be added to the index.
      // we need to setup the crs and check that settings like cut before adding are compatible with
      // the actual geometry

      // TODO add a consistency check?
      srid = geom.getSRID();
      // NB this is going to throw an exception if the geometry has no valid srid - but that's OK, this index needs to
      // be referenced to work properly
      crs = sridSet.get(srid);
      indexedFamily = GeometryFamily.fromClass(geom.getClass());

      if (options.cutBeforeAdding.isEmpty()) {
        // cutBeforeAdding will default to true, but only for polygons as we don't support cutting
        // other geometry types.
        boolean cutBeforeAdding = indexedFamily == GeometryFamily.POLYGONAL;
        log.debug("cutBeforeAdding default value is set to {}", cutBeforeAdding);
        options.cutBeforeAdding = Optional.of(cutBeforeAdding);
      }

      // unsupported - seems like a bit of an edge case and I don't want to test it and make it work at this moment.
      // A quick test using an existing model that did sample lines gave inconsistent results to the non cutting
      // version, though that might have been because the model wasn't deterministic
      if (options.cutBeforeAdding.get() && indexedFamily != GeometryFamily.POLYGONAL) {
        throw new RiskscapeException(
            GeneralProblems.get().operationNotSupported("intersection index cutting + " + indexedFamily, getClass())
                .withChildren(PROBLEMS.intersectionCutNotSupportedHint(indexedFamily))
        );
      }
    }

    Envelope geomEnvelope = geom.getEnvelopeInternal();
    if (options.cutBeforeAdding.get()) {
      if (!indexedFamily.isSameFamily(geom)) {
        // the current geometry is not of the same family that the index was set up with. we must be dealing with
        // a mixed bag of geometry types.
        // this is a problem when cutting because cutting is only supported for polygons.
        GeometryFamily currentGeomFamily = GeometryFamily.fromClass(geom.getClass());
        throw new RiskscapeException(
            GeneralProblems.get().operationNotSupported("intersection index cutting + " + currentGeomFamily, getClass())
                .withChildren(PROBLEMS.intersectionCutNotSupportedHint(currentGeomFamily))
        );
      }
      if (options.isTooComplex(geom, geomEnvelope)) {
        cutComplexGeometry(tuple, geom, geomEnvelope);
      } else {
        tree.insert(geomEnvelope, Pair.of(tuple, geom));
      }
    } else {
      tree.insert(geomEnvelope, tuple);
    }
    // regardless of cutting, the envelope is expanded by the feature's envelope - no need to do it per cut (will still
    // be the same bounds)
    envelope.expandToInclude(geomEnvelope);
  }

  /**
   * Recursively cuts the given geometry in to quads until its constituent parts are no longer "too complex"
   */
  private void cutComplexGeometry(Tuple tuple, Geometry geom, Envelope env) {
    LinkedList<Geometry> stack = new LinkedList<>();
    // we already know this is too complex - we've tested it before getting here
    cutOnce(geom, env, stack);

    while (!stack.isEmpty()) {
      Geometry candidate = stack.removeFirst();
      Envelope candidateEnv = candidate.getEnvelopeInternal();

      // if the given thing is too complex, queue it up to be broken apart, or
      if (options.isTooComplex(candidate, candidateEnv)) {
        cutOnce(candidate, candidateEnv, stack);
      } else {
        tree.insert(candidateEnv, Pair.of(tuple, candidate));
      }
    }
  }

  /**
   * Cut the given geometry in to quads, adding the new geometries to the stack
   */
  private void cutOnce(Geometry geom, Envelope env, LinkedList<Geometry> stack) {
    Collection<Envelope> envelopes = Quadrant.partition(env).values();

    for (Envelope quadEnvelope : envelopes) {
      LinearRing envRing = RecursiveQuadGridOp.envelopeToRing(geom.getFactory(), quadEnvelope);
      Geometry intersection = geom.intersection(new Polygon(envRing, null, envRing.getFactory()));

      // NB this creates geometries of different types when intersections are 'glancing'.  We could possibly filter
      // those out, but I would imagine it'll make little difference either way
      if (intersection.isEmpty()) {
        continue;
      }

      stack.add(intersection);
    }
  }


  public ReferencedEnvelope getReferencedEnvelope() {
    return new ReferencedEnvelope(envelope.getMinX(), envelope.getMaxX(), envelope.getMinY(), envelope.getMaxY(), crs);
  }

  /**
   * Query the index for any tuples that intersect the given geometry, returning the intersecting geometry and the
   * tuple.
   *
   * Discards any intersections that are not of an 'equivalent' geometry type as the product of intersecting the
   * given geometry type and the indexed geometry type, according to
   * {@link GeometryFamily}.  These intersections are typically caused by glancing intersections where features share
   * borders and vertices
   * @param geom     the geomerty to find the intersections on
   * @return Tuple   intersections the parameter geometry
   */
  public List<Pair<Geometry, Tuple>> findIntersections(Geometry geom) {
    return findIntersections(geom, g -> {});
  }

  /**
   * @return lookupGeom reprojected to match the IntersectionIndex contents
   */
  private Geometry reproject(Geometry lookupGeom) {
    if (tree.isEmpty()) {
      // tree is empty so it doesn't matter what CRS the lookup geometry is in. Otherwise srid == -1,
      // and we'll die trying to reproject - see #1405
      return lookupGeom;
    }

    // SridSet will only reproject if that is necessary
    return sridSet.reproject(lookupGeom, srid);
  }

  /**
   * Find tuples that contain a geometry that intersects with the given geometry.
   *
   * Similar to {@link #findIntersections(org.locationtech.jts.geom.Geometry) } except that:
   * - intersecting geometries are not returned
   * - will include values from geometries that have glancing intersections that are excluded from findIntersections
   *
   * @param geom geometry to find intersections for
   * @return list of tuples that have some intersection with geom.
   */
  @SuppressWarnings("unchecked")
  public List<Tuple> findIntersectingValues(Geometry geom) {
    // reproject the input geometry (if necessary) to match the index's CRS
    Geometry reprojected = reproject(geom);

    return tuplesInEnvelope(reprojected.getEnvelopeInternal())
        // we need to check that the geometries intersect. this should be a little faster than computing
        // the intersection as no new geometries need to be created.
        .filter(t -> {
          Geometry toMatch = geometryMemberAccessor.evaluate(t, Geometry.class);
          return toMatch.intersects(reprojected);
        })
        .toList();
  }

  /**
   * Similar to {@link #findIntersections(org.locationtech.jts.geom.Geometry) } but also returns an
   * optional geometry containing the difference. That is the area of geom that did not intersect
   * with any features from the index. This will be empty if there is no difference.
   *
   * This method will calculate the difference in the CRS of the indexed geometries. This is to remove
   * rounding errors that would occur if the difference was taken from geometries that had already been
   * re-projected to another CRS.
   *
   * This method should only be used if the difference is required. Calculating the difference could
   * be expensive, so if it is not needed use {@link #findIntersections(org.locationtech.jts.geom.Geometry) }.
   *
   * @param geom   the geomerty to find the intersections on.
   * @return Geometry, containing the area of geom that did not intersect, and intersections the
   * parameter geometry.
   */
  public Pair<Optional<Geometry>, List<Pair<Geometry, Tuple>>> findDifferenceAndIntersections(Geometry geom) {
    List<Geometry> unprojectedMatches = Lists.newArrayList();
    List<Pair<Geometry, Tuple>> intersections = findIntersections(geom, g -> unprojectedMatches.add(g));

    if (intersections.isEmpty()) {
      // If there are no intersections then we know that the difference *is* the input geom.
      // We can short circuit a lot of work now and return geom as the difference. This also removes
      // any reprojection type noise.
      return Pair.of(Optional.of(geom), intersections);
    }

    // reproject the input geometry (if necessary) to match the index's CRS
    Geometry reprojected = reproject(geom);

    // combine all the potentially overlapping geometry into a single shape. If the index is an
    // area-layer, then any neighbouring regions will naturally merge together into a single polygon.
    // If we dealt with each shape individually then we can end up with reprojection noise between the
    // area boundaries where we end up with tiny gaps. We avoid that with this approach
    Geometry combined = reprojected.getFactory().createPoint();
    for (Geometry unprojectedOverlap: unprojectedMatches) {
      combined = combined.union(unprojectedOverlap);
    }

    // find any geometry that falls outside of the combined blob
    GeometryFamily reprojectedGeomFamily = GeometryFamily.from(reprojected);
    Geometry remainder = GeometryUtils.removeNonFamilyMembers(
        overlayOperations.difference(reprojected.copy(), combined),
        reprojectedGeomFamily);

    if (remainder.isEmpty()) {
      return Pair.of(Optional.empty(), intersections);
    } else {
      // reproject the difference back to the original CRS
      return Pair.of(Optional.of(sridSet.reproject(remainder, geom.getSRID())), intersections);
    }
  }

  /**
   * @param geom          geometry to find intersections for
   * @param matchConsumer a consumer that will accept any matching geometry in the index that overlaps with the
   *                      given geometry. Note that the original/unmodified geometry is used (not the intersection)
   *                      in the index's CRS (which may be different to the input geometry).
   * @return list of pairs of intersecting geometry and tuple. The intersecting geometry is reprojected
   *         (if necessary) to the CRS of the input geom.
   */
  @SuppressWarnings("unchecked")
  private List<Pair<Geometry, Tuple>> findIntersections(Geometry geom, Consumer<Geometry> matchConsumer) {
    // reproject the input geometry (if necessary) to match the index's CRS
    Geometry reprojected = reproject(geom);

    List<Tuple> found = tuplesInEnvelope(reprojected.getEnvelopeInternal()).toList();
    List<Pair<Geometry, Tuple>> intersecting = new ArrayList<>(found.size());

    // work out the geometry we expect back from the intersection
    // - this protects against a polygon becoming a line/point because it touches another polygon
    // - ditto for line becoming a point
    // - but still allowing a polygon to intersect a point/line relation and result in those types
    GeometryFamily expectedType = GeometryFamily.fromClass(geom.getClass()).min(indexedFamily);

    for (Tuple t : found) {
      Geometry toMatch = geometryMemberAccessor.evaluate(t, Geometry.class);

      // The intersection should be ignored when:
      // - it is empty
      // - it is a different kind of geometry family to the inputs (and thus becomes empty)
      Geometry intersection = GeometryUtils.removeNonFamilyMembers(
          overlayOperations.intersection(toMatch, reprojected), expectedType);

      if (intersection.isEmpty()) {
        continue;
      }

      matchConsumer.accept(toMatch);

      if (geom.getSRID() != intersection.getSRID() && intersection.equalsTopo(reprojected)) {
        // if the intersection is topology equal to reprojected then the input geom is completely covered
        // by the indexed geom. In this case we can return geom (rather than a reprojected (x2) version
        // of it. This removes the reprojection rounding errors. But we don't want to waste time on
        // topo equals if we are in the same projection anyway.
        intersecting.add(new Pair<>(geom, t));
      } else {
        // reproject back to origin CRS if necessary
        intersecting.add(new Pair<>(sridSet.reproject(intersection, geom.getSRID()), t));
      }
    }

    return intersecting;
  }

  @SuppressWarnings("unchecked")
  private Stream<Tuple> tuplesInEnvelope(Envelope env) {
    if (options.cutBeforeAdding.get()) {
      // in cut mode, we store pairs
      List<Pair<Tuple, Geometry>> list = (List<Pair<Tuple, Geometry>>) tree.query(env);
      return list.stream()
          .map((pair) -> pair.getLeft())
          // because we've cut each feature up multiple times, we're potentially going to get many hits on the same
          // tuple, we only keep one of each, as we will compute the intersection as a whole.
          // NB we could take a different approach and instead compute the intersection against each chunk and then
          // 'dissolve' those features, but that would need a bit of testing both functional and performant to see if
          // it's worth it.
          // For now, the main use case for cutting is to speed up point intersections to 'find' the best feature
          .distinct();
    }
    return tree.query(env).stream();
  }

  /**
   * Query the index for any tuples that intersect with the given point, handling any required crs differences
   */
  public List<Tuple> findPointIntersections(Point point) {
    Geometry reprojected = reproject(point);

    List<Tuple> intersections;
    if (options.cutBeforeAdding.get()) {
      @SuppressWarnings("unchecked")
      List<Pair<Tuple, Geometry>> queried = tree.query(reprojected.getEnvelopeInternal());
      intersections = new ArrayList<>(queried.size());

      for (Pair<Tuple, Geometry> pair : queried) {
        Geometry indexedGeom = pair.getRight();
        Tuple tuple = pair.getLeft();

        if (indexedGeom.intersects(reprojected)) {
          intersections.add(tuple);
        }
      }
    } else {
      @SuppressWarnings("unchecked")
      List<Tuple> queried = tree.query(reprojected.getEnvelopeInternal());
      intersections = new ArrayList<>(queried.size());

      for (Tuple tuple : queried) {
        Geometry indexedGeom = (Geometry) geometryMemberAccessor.evaluate(tuple);

        if (indexedGeom.intersects(reprojected)) {
          intersections.add(tuple);
        }
      }
    }

    return intersections;
  }

  /**
   * @return the size of the index. Only intended for test use hence package private.
   */
  int size() {
    return tree.size();
  }

}
