/*
 * 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 org.geotools.api.referencing.crs.CoordinateReferenceSystem;
import org.locationtech.jts.geom.Coordinate;
import org.locationtech.jts.geom.Envelope;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.index.kdtree.KdNode;
import org.locationtech.jts.index.kdtree.KdNodeVisitor;
import org.locationtech.jts.index.kdtree.KdTree;

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.expr.StructMemberAccessExpression;
import nz.org.riskscape.engine.types.Struct.StructMember;

/**
 * An index that when queried will return the nearest neighbour should one exist within a max distance.
 */
@Slf4j
public class NearestNeighbourIndex {

  @RequiredArgsConstructor
  public static class Visitor implements KdNodeVisitor {

    public final Coordinate lookFor;
    public KdNode closest;
    public double lastDistance = 0;
    public double currentDistance = Double.POSITIVE_INFINITY;
    public int visitCount = 0;
    public boolean ignoreRadiusCompensation = false;

    @Override
    public void visit(KdNode node) {
      double distance = node.getCoordinate().distance(lookFor);
      visitCount++;

      if (distance > currentDistance && !ignoreRadiusCompensation) {
        // this is outside of the radius of the circle that fits within our rectangle, so we can't tell if this is
        // nearest or not
        return;
      }
      if (closest == null) {
        closest = node;
        lastDistance = distance;
      } else if (distance < lastDistance) {
        closest = node;
        lastDistance = distance;
      }
    }
  }

  private final Envelope envelope = new Envelope();

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

  /**
   * The initial search size to start with.  Set this to twice the avg distance between points for optimal
   * performance.
   *
   * This size is used to create a search envelope for querying the kdtree.  If the search finds no results, then
   * the envelope is doubled in size.  The search continues until the envelope grows to cover the whole search area.
   */
  private final double startSize;

  @Getter
  private final double maxDistanceInCrsUnits;

  private final KdTree tree = new KdTree();

  /**
   * All geometry must have this srid or adding tuples to the index will fail
   */
  private final int expectedSrid;

  /**
   * Create a nearest neighbour index with the given max distance (in metres).
   * @param geometryMember    the geometry member that contains the geometry we are indexing from given tuples
   * @param sridSet an srid set to use for asserting that the geometry we end up indexing is in the correct CRS.
   * @param crs the coordinate reference system this index will be in.  So far, this is not persisted with the index,
   * but is used to derive a metric max distance and an expectedSrid
   * @param maxDistance the max distance (in metres ) that the match must be within.  Note that if the CRS is not based
   * on a metric grid, this will be an approximation.
   */
  public static NearestNeighbourIndex metricMaxDistance(
      StructMember geometryMember,
      SRIDSet sridSet,
      CoordinateReferenceSystem crs,
      double maxDistance) {
    int expectedSrid = sridSet.get(crs);
    double distanceInCrsUnits = GeometryUtils.toCrsUnits(maxDistance, crs);

    return new NearestNeighbourIndex(geometryMember, distanceInCrsUnits, expectedSrid);
  }

  // private - we want the static constructor with the illustrative name to be used
  private NearestNeighbourIndex(StructMember geometryMember, Double maxDistanceInCrsUnits, int expectedSrid) {
    this.expectedSrid = expectedSrid;
    this.geometryMemberAccessor = new StructMemberAccessExpression(false, geometryMember);
    this.maxDistanceInCrsUnits = maxDistanceInCrsUnits;
    this.startSize = maxDistanceInCrsUnits * 2;
  }

  /**
   * Insert a {@link Tuple} into the index. The tuple will be indexed by to the geometry accessed with
   * {@link #geometryMemberAccessor}.
   *
   * @param item to add to index
   */
  public final void insert(Tuple item) {
    Geometry geom = geometryMemberAccessor.evaluate(item, Geometry.class);
    if (geom == null) {
      // null geometries are not indexed
      return;
    }

    // sanity check
    if (geom.getSRID() != expectedSrid) {
      throw new AssertionError("NearestNeighbourIndex does not support mixed CRS - expected SRID %d, got %d"
          .formatted(expectedSrid, geom.getSRID()));
    }

    tree.insert(geom.getCentroid().getCoordinate(), item);
    envelope.expandToInclude(geom.getEnvelopeInternal());
  }

  /**
   * @return the envelope that the index may provide matches within.
   */
  public Envelope getEnvelope() {
    // Envelope is the bounds of the indexed points. So we expand it by maxDistance to made the bounds that
    // this index may provide a valid answer for.
    return new Envelope(envelope.getMinX() - maxDistanceInCrsUnits,
        envelope.getMaxX() + maxDistanceInCrsUnits,
        envelope.getMinY() - maxDistanceInCrsUnits,
        envelope.getMaxY() + maxDistanceInCrsUnits);
  }

  /**
   * Find the indexed entry that is closest to coordinate but still within {@link #maxDistanceInCrsUnits}.
   * @param coordinate where to query the index
   * @return closest entry to coordinate, or null if there no entry is found with maxDistance.
   */
  public Tuple query(Coordinate coordinate) {
    if (tree.isEmpty()) {
      return null;
    }

    Envelope bounds = getEnvelope();

    double currentSize = startSize / 2;
    Visitor visitor = new Visitor(coordinate);
    Envelope initialEnv = new Envelope(coordinate);
    Envelope queryEnv = new Envelope(coordinate);
    int expansions = 0;

    while (!queryEnv.covers(bounds)) {
      queryEnv.init(initialEnv);
      queryEnv.expandBy(currentSize);
      visitor.currentDistance = currentSize;

      if (queryEnv.covers(bounds)) {
        // this is searching the whole index now - tell the visitor to ignore distance
        visitor.ignoreRadiusCompensation = true;
      }

      tree.query(queryEnv, visitor);

      if (visitor.closest != null) {
        if (visitor.visitCount > 100) {
          log.info("Item found after {} visits, consider decreasing start bounds size ", visitor.visitCount);
        }
        Coordinate foundAt = visitor.closest.getCoordinate();
        if (foundAt.distance(coordinate) > maxDistanceInCrsUnits) {
          // sadly it is too far away
          break;
        }
        return (Tuple)visitor.closest.getData();
      }
      if (currentSize > maxDistanceInCrsUnits) {
        // current size is already too big and we haven't found anything, so let's bail
        break;
      }

      currentSize *= 2;
      expansions++;

      if (expansions > 5) {
        log.warn("Nothing found within envelope after 5 expansions {}...", expansions);
      }
    }

    return null;
  }

}
