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

import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.BiFunction;
import java.util.function.Function;

import org.geotools.geometry.jts.ReferencedEnvelope;
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 org.geotools.api.referencing.crs.CoordinateReferenceSystem;

import com.google.common.collect.Lists;

import lombok.AccessLevel;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.SRIDSet;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.coverage.TypedCoverage;
import nz.org.riskscape.engine.data.coverage.SpatialRelationTypedCoverage;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Struct.StructMember;
import nz.org.riskscape.engine.util.Pair;

@RequiredArgsConstructor(access = AccessLevel.PRIVATE)
public class MultiCoverage implements TypedCoverage {

  /**
   * Functional interface for an object that can set raw values with the result of evaluating a
   * coverage.
   */
  @FunctionalInterface
  public interface RawValueSetter {

    /**
     * Set the value (or parts of it) into rawValues in the appropriate positions for {@link MultiCoverage#type}.
     *
     * @param value     to set into the appropriate parts of rawValues
     * @param tuple store the results into this
     */
    void set(Object value, Tuple tuple);
  }

  /**
   * Creates a new MultiCoverage that yields the given struct `type` by sampling all of the coverages and setting those
   * to type.  Each member of the type must correspond to a coverage.
   * @param type          the type to be returned from the coverage
   * @param coverages     the coverages that make up the multi coverage
   * @param valueSetters  to evaluate the coverage and set the results in the appropriate positions
   *                           of the tuples raw results (as given by #type)
   * @param cutFunction   a function that can cut a given geometry into smaller pieces for sampling the coverages.
   * Used by {@link #getEvaluateIntersectionOp()}.
   */
  public static MultiCoverage create(
      Struct type,
      TypedCoverage[] coverages,
      RawValueSetter[] valueSetters,
      SRIDSet sridSet,
      BiFunction<Geometry, Point, List<Geometry>> cutFunction
  ) {
    if (coverages.length == 0) {
      throw new IllegalArgumentException("can not pass an empty map");
    }

    if (coverages.length != valueSetters.length) {
      throw new IllegalArgumentException("must have same number of coverages and evaluators");
    }

    return new MultiCoverage(
      type,
      coverages,
      valueSetters,
      sridSet,
      cutFunction
    );
  }

  @Getter
  private final Struct type;

  private final TypedCoverage[] coverages;

  private final RawValueSetter[] valueSetters;

  private Envelope envelope;

  private final SRIDSet sridSet;

  private final BiFunction<Geometry, Point, List<Geometry>> cutFunction;

  @Override
  public Object evaluate(Point point) {
    Tuple results = new Tuple(type);
    boolean anyResult = false;

    for (int i = 0; i < coverages.length; i++) {
      TypedCoverage coverage = coverages[i];
      RawValueSetter valueSetter = valueSetters[i];

      Object evaluated = coverage.evaluate(point);
      if (evaluated != null) {
       valueSetter.set(evaluated, results);
       anyResult = true;
      }
    }

    // NB: sampling a coverage returns null when there's no data there. For a
    // multi-coverage, we should still return null if all coverages were null
    return anyResult ? results : null;
  }

  /**
   * This samples each of the coverages the MultiCoverage contains individually, and then
   * returns the combined result. When evaluating intersections, the input geometry is cut into
   * pieces (i.e. grid segments), and the *centre* of each piece is sampled individually.
   * It is important to note that we only ever use {@link TypedCoverage#evaluate(Point)} to
   * sample each individual coverage, even if the coverage is a {@link SpatialRelationTypedCoverage}.
   */
  @Override
  public Optional<Function<Geometry, List<Pair<Geometry, Object>>>> getEvaluateIntersectionOp() {
    return Optional.of(this::evaluateIntersection);
  }

  List<Pair<Geometry, Object>> evaluateIntersection(Geometry geom) {

    // for best sampling results, we want to align the geometry to a common grid origin
    // when cutting it. Align everything to the grid bounds of the first coverage,
    // and if that's not known then fallback to using the geometry's bounds
    Point origin = coverages[0].getEnvelope()
        .map(e -> getGridOrigin(e, coverages[0].getCoordinateReferenceSystem()))
        .orElse(getGridOrigin(geom.getEnvelopeInternal(), sridSet.get(geom)));

    // cut the input geometry so it better matches the underlying coverages
    // (ideally this would cut it by the same raster grid that our coverages use)
    List<Geometry> griddedGeom = cutFunction.apply(geom, origin);

    List<Pair<Geometry, Object>> results = new LinkedList<>();
    for (Geometry gridSegment : griddedGeom) {
      Point point = getCentre(gridSegment);
      Tuple tuple = new Tuple(type);
      boolean anyResult = false;

      for (int i = 0; i < coverages.length; i++) {
        TypedCoverage coverage = coverages[i];
        RawValueSetter valueSetter = valueSetters[i];

        Object evaluated = coverage.evaluate(point);
        if (evaluated != null) {
          valueSetter.set(evaluated, tuple);
          anyResult = true;
        }
      }
      // if sampling all the coverages did not produce any result, we then don't
      // want to add the tuple to the list (it's the 'no data' case)
      if (anyResult) {
        results.add(Pair.of(gridSegment, tuple));
      }
    }
    return results;
  }

  @Override
  public Optional<Function<Geometry, List<? extends Object>>> getEvaluateIntersectingValuesOp() {
    // for multi coverage we just defer to getEvaluateIntersectionOp() and strip out the geometries we
    // don't want. we do this because apart from including the the geometries getEvaluateIntersectionOp
    // doesn't do anything else that we could skip for better performance anyway.
    return getEvaluateIntersectionOp().map(op ->
        op.andThen(list ->
            Lists.transform(list, pair -> pair.getRight())
        )
    );
  }

  @Override
  public CoordinateReferenceSystem getCoordinateReferenceSystem() {
    return coverages[0].getCoordinateReferenceSystem();
  }

  @Override
  public Optional<ReferencedEnvelope> getEnvelope() {
     if (envelope == null) {
      envelope = coverages[0].getEnvelope().orElse(null);

      if (envelope == null) {
        return Optional.empty();
      }

      // avoid any state modifying bugs
      envelope = envelope.copy();

      for (int i = 1; i < coverages.length; i++) {
        TypedCoverage typedCoverage = coverages[i];
        ReferencedEnvelope expandWith = typedCoverage.getEnvelope().orElse(null);
        // TODO how to accurately reproject a bounding box?  Just translating tl and br won't work well enough...
        if (expandWith != null && typedCoverage.getCoordinateReferenceSystem().equals(getCoordinateReferenceSystem())) {
          envelope.expandToInclude(expandWith);
        }
      }
    }
    return Optional.of(new ReferencedEnvelope(envelope, getCoordinateReferenceSystem()));
  }

  @Override
  public String toString() {
    return String.format("MultiCoverage(type=%s)", getType());
  }

  public Map<String, TypedCoverage> getCoveragesAsMap() {
    LinkedHashMap<String, TypedCoverage> map = new LinkedHashMap<>();
    List<StructMember> keys = type.getMembers();
    for (int i = 0; i < coverages.length; i++) {
      map.put(keys.get(i).getKey(), coverages[i]);
    }
    return map;
  }

  private Point getGridOrigin(Envelope env, CoordinateReferenceSystem crs) {
    Coordinate origin = new Coordinate(env.getMinX(), env.getMinY());
    return sridSet.getGeometryFactory(crs).createPoint(origin);
  }

  private Point getCentre(Geometry geom) {
    return geom.getFactory().createPoint(geom.getEnvelopeInternal().centre());
  }
}
