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

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

import com.google.common.collect.BoundType;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Range;

import lombok.Data;
import lombok.Getter;
import lombok.NonNull;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.types.Types;

/**
 * An {@link RiskscapeFunction} that allows a discrete function to be defined.  Can be defined from a mix of constant
 * values and other {@link RiskscapeFunction}s points, ranges and points.
 */
public class DiscreteFunction implements RiskscapeFunction {

  /**
   * API For users to build their {@link DiscreteFunction}
   *
   * ```
   * // 1     :    5
   * // 2..10 :    5x + 1
   * // - plus 1..2 gets a line from 5 to 11
   * builder()
   *   .add(1, 5)
   *   .add(2, 10, Maths.newPolynomial(1, 5))
   *   .withLinearInterpolation()
   *   .build()
   * ```
   */
  public static class Builder {

    private List<Pair> soFar = new ArrayList<>();
    private boolean linearInterpolation = false;
    private Type argumentType = Types.INTEGER;
    private Type returnType = Types.INTEGER;
    private boolean closeUpperBound = true;

    /**
     * @return a {@link DiscreteFunction} based on the construction so far.
     */
    public DiscreteFunction build() {

      // I guess?
      if (soFar.isEmpty()) {
        return new DiscreteFunction(soFar, argumentType, returnType);
      }

      if (linearInterpolation) {
        List<Pair> toAdd = new ArrayList<>();

        fillGaps(toAdd);

        for (Pair gap : toAdd) {
          add(gap);
        }
      }

      if (closeUpperBound) {
        closeDisconnectedUpperBounds();
      }

      return new DiscreteFunction(soFar, argumentType, returnType);
    }

    /**
     * Add a function, applicable for the range (lower >= x > upper)
     * @param lower inclusive start of the x range
     * @param upper exclusive end of the x range
     * @param function applicable function for the range
     * @return the builder, to allow a fluid API
     * @throws IllegalArgumentException if the range overlaps, or if any of the function types are not Numeric
     */
    public Builder addFunction(double lower, double upper, RiskscapeFunction function) {
      return add(Range.closedOpen(lower, upper), function);
    }

    /**
     * Add a constant, applicable for the range (lower >= x > upper)
     * @param lower inclusive start of the x range
     * @param upper exclusive end of the x range
     * @param constant value to return for the range
     * @return the builder, to allow a fluid API
     * @throws IllegalArgumentException if the range overlaps, or if any of the function types are not Numeric
     */
    public Builder addConstant(double lower, double upper, Number constant) {
      return add(Range.closedOpen(lower, upper), Maths.newConstant(constant));
    }

    /**
     * Add a function applicable only for the given value
     * @param at single xvalue that the function is applicable for
     * @param function applicable function for the range
     * @return the builder, to allow a fluid API
     * @throws IllegalArgumentException if the range overlaps, or if any of the function types are not Numeric
     */
    public Builder addPoint(double at, RiskscapeFunction function) {
      return add(Range.singleton(at), function);
    }

    /**
     * Add a constant value, applicable only for the given at value
     * @param at single x value that the function is applicable for
     * @param constant the value to return
     * @return the builder, to allow a fluid API
     * @throws IllegalArgumentException if the range overlaps, or if any of the function types are not Numeric
     */
    public Builder addPoint(double at, Number constant) {
      return add(Range.singleton(at), Maths.newConstant(constant));
    }

    /**
     * Add a function applicable for the given range
     * @param range any valid {@link Range}
     * @param function applicable function for the range
     * @return the builder, to allow a fluid API
     * @throws IllegalArgumentException if the range overlaps, or if any of the function types are not Numeric
     */
    public Builder add(Range<Double> range, RiskscapeFunction function) {
      return add(new Pair(range, function));
    }

    /**
     * Add an already constructed {@link Pair} to the builder
     * @return the builder, to allow a fluid API
     * @throws IllegalArgumentException if the range overlaps, or if any of the function types are not Numeric
     */
    public Builder add(Pair pair) {

      List<Type> args = pair.getFunction().getArgumentTypes();

      if (args.size() != 1) {
        throw new IllegalArgumentException(String.format(
            "Functions must declare exactly 1 argument, not %d - %s", args.size(), pair));
      }

      this.returnType = highestRanking(this.returnType, pair.function.getReturnType());
      this.argumentType = highestRanking(this.argumentType, args.get(0));

      for (Pair checkAgainst : soFar) {

        if (!checkAgainst.range.isConnected(pair.range)) {
          continue;
        }

        Range<Double> intersection = checkAgainst.getRange().intersection(pair.range);
        if (!intersection.isEmpty()) {
          throw new IllegalArgumentException(String.format(
              "Functions have overlapping ranges. %s overlaps %s", pair, checkAgainst));
        }
      }

      soFar.add(pair);
      return this;
    }

    public Builder withoutUpperBoundClosing() {
      this.closeUpperBound = false;
      return this;
    }

    /**
     * Fill in any gaps between ranges with a linear function that bridges the gap.
     */
    public Builder withLinearInterpolation() {
      linearInterpolation = true;
      return this;
    }

    /**
     * Return either lhs or rhs, depending on which is higher according to {@link DiscreteFunction#NUMERIC_TYPE_RANK}.
     * This means that the returned type will be capable of representing either's values
     */
    private Type highestRanking(Type lhs, Type rhs) {
      if (NUMERIC_TYPE_RANK.indexOf(lhs) > NUMERIC_TYPE_RANK.indexOf(rhs)) {
        return lhs;
      } else {
        return rhs;
      }
    }


    /**
     * Find gaps between bounded ranges, fill them with a linear function from x1 to x2
     */
    private void fillGaps(List<Pair> gapFills) {
      // work from right to left, looking for gaps
      // PS nothing overlaps, we've asserted that in the #add method
      for (Pair lower : soFar) {
        if (lower.range.hasUpperBound()) {

          Pair next = null;
          double lowersUpper = lower.range.upperEndpoint();
          double min = Double.MAX_VALUE;
          for (Pair pair : soFar) {
            if (lower == pair || !pair.getRange().hasLowerBound()) {
              continue;
            }
            double pairsLower = pair.range.lowerEndpoint();

            if (pairsLower > lowersUpper && pairsLower < min) {
              next = pair;
              min = pairsLower;
            }
          }

          if (next != null) {
            Range<Double> gap = Range.range(
                Double.valueOf(lowersUpper),
                flip(lower.range.upperBoundType()),
                Double.valueOf(min),
                flip(next.range.lowerBoundType()));

            Double x1 = (Double) Types.FLOATING.coerce(gap.lowerEndpoint());
            Double x2 = (Double) Types.FLOATING.coerce(gap.upperEndpoint());

            Double y1 = (Double) Types.FLOATING.coerce(lower.function.call(Collections.singletonList(x1)));
            Double y2 = (Double) Types.FLOATING.coerce(next.function.call(Collections.singletonList(x2)));

            double slope;
            if (y2.equals(y1)) {
              slope = 0;
            } else {
              slope = (y2 - y1) / (x2.doubleValue() - x1.doubleValue());
            }
            double c = y1 - (slope * x1.doubleValue());
            gapFills.add(new Pair(gap, Maths.newPolynomial(new double[] {c, slope})));
          }
        }
      }
    }

    private void closeDisconnectedUpperBounds() {

      // find a range with an open upper bound
      // see if it is connected to another range
      // if not, close it
      List<Pair> toRemove = new ArrayList<>();
      List<Pair> toAdd = new ArrayList<>();

      for (Pair pair : soFar) {
        Range<Double> range = pair.range;
        if (!range.hasUpperBound() || range.upperBoundType() == BoundType.CLOSED) {
          continue;
        }

        Pair neighbour = null;
        for (Pair possibleNeighbour : soFar) {
          // it's not likely, but there could be a range with an open lower bound, which is not considered connected
          // - if someone is adding Ranges using the Range API, we should trust that they know what they are doing
          if (!possibleNeighbour.range.hasLowerBound() || possibleNeighbour.range.lowerBoundType() == BoundType.OPEN) {
            continue;
          }

          if (possibleNeighbour.range.lowerEndpoint().equals(range.upperEndpoint())) {
            neighbour = possibleNeighbour;
            break;
          }
        }

        // no neighbour found, close it up
        if (neighbour == null) {
          Range<Double> newRange;
          if (range.hasLowerBound()) {
            newRange = Range.range(
                range.lowerEndpoint(), range.lowerBoundType(),
                range.upperEndpoint(), BoundType.CLOSED);
          } else {
            newRange = Range.upTo(range.upperEndpoint(), BoundType.CLOSED);
          }

          // swapsies
          toRemove.add(pair);
          toAdd.add(new Pair(newRange, pair.function));
        }
      }

      for (Pair remove : toRemove) {
        soFar.remove(remove);
      }

      for (Pair add : toAdd) {
        // use the add method so we get validation of overlapping ranges, just in case
        add(add);
      }

    }

    /**
     * Reverse the given bound for the range
     */
    private BoundType flip(BoundType boundType) {
      if (boundType == BoundType.CLOSED) {
        return BoundType.OPEN;
      } else {
        return BoundType.CLOSED;
      }
    }
  }

  /**
   * @return a new {@link Builder} for constructing a {@link DiscreteFunction}.
   */
  public static Builder builder() {
    return new DiscreteFunction.Builder();
  }

  /**
   * A member of the {@link DiscreteFunction}.  Links a {@link Range} to an {@link RiskscapeFunction}.
   */
  @Data
  public static class Pair {
    private final Range<Double> range;
    private final RiskscapeFunction function;
  }

  /**
   * Ranks numeric {@link Type}s, such that a higher ranked {@link Type} can represent any of the lower ranked
   * {@link Type}s.
   * TODO consider moving this to the NumberTypes enum
   */
  public static final List<Type> NUMERIC_TYPE_RANK = ImmutableList.of(
    Types.INTEGER,
    Types.FLOATING,
    Types.DECIMAL
  );

  @Getter
  private final List<Pair> pairs;

  @Getter
  private final List<Type> argumentTypes;

  @Getter
  private final Type returnType;

  private DiscreteFunction(List<Pair> functions, Type argumentType, Type returnType) {
    this.pairs = functions;
    this.argumentTypes = Collections.singletonList(argumentType);
    this.returnType = returnType;
  }

  @Override
  public Object call(@NonNull List<Object> args) {
    if (args.size() != 1) {
      throw new IllegalArgumentException(
          String.format("this function requires exactly 1 argument, %d given", args.size()));
    }

    Double arg = (Double) Types.FLOATING.coerce(args.get(0));

    for (Pair pair : pairs) {
      if (pair.getRange().contains(arg)) {
        RiskscapeFunction function = pair.getFunction();
        Object coercedArg = function.getArgumentTypes().get(0).coerce(arg);
        Object rawReturned = function.call(Collections.singletonList(coercedArg));

        return getReturnType().coerce(rawReturned);
      }
    }

    throw new IllegalArgumentException(String.format("%s is outside of the functions range - %s",
        args.get(0),
        pairs.stream()
          .map(Pair::getRange)
          .map(Range::toString)
          .collect(Collectors.joining(", "))
    ));
  }

}
