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

import java.util.List;
import java.util.Optional;
import java.util.function.BiFunction;
import java.util.stream.Collectors;

import lombok.Getter;
import lombok.RequiredArgsConstructor;
import nz.org.riskscape.engine.ArgsProblems;
import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.function.ArgumentList;
import nz.org.riskscape.engine.function.BaseMathsFunction;
import nz.org.riskscape.engine.function.IdentifiedFunction;
import nz.org.riskscape.engine.function.RiskscapeFunction;
import nz.org.riskscape.engine.rl.RealizableFunction;
import nz.org.riskscape.engine.rl.RealizationContext;
import nz.org.riskscape.engine.types.Nullable;
import nz.org.riskscape.engine.types.Struct;
import nz.org.riskscape.engine.types.Struct.StructMember;
import nz.org.riskscape.engine.types.Type;
import nz.org.riskscape.engine.types.Types;
import nz.org.riskscape.problem.ResultOrProblems;
import nz.org.riskscape.rl.ast.FunctionCall;

/**
 * Function that will return a scaled version of the first argument by the scale factor in the
 * second argument.
 *
 * This function is only applicable to numeric {@link Types}s. If the first argument is not numeric or is a
 * {@link Struct} with no numeric members then the function will not be realized.
 *
 * In the case of {@link Struct} inputs only numeric members are scaled. Other non-numeric members
 * are left as is. This included any nested {@link Struct} e.g. any nested structs will not be scaled at all.
 */
public class Scale extends BaseMathsFunction implements IdentifiedFunction, RealizableFunction {

  @RequiredArgsConstructor
  private static class ScaleImpl implements RiskscapeFunction {

    @Getter
    private final List<Type> argumentTypes;
    @Getter
    private final Type returnType;
    private final BiFunction<Object, Double, Object> scaleFunction;

    @Override
    public Object call(List<Object> args) {
      Double scaleFactor = (Double)args.get(1);
      return scaleFunction.apply(args.get(0), scaleFactor);
    }

  }

  @Getter
  private final ArgumentList arguments = ArgumentList.create(
      "input_value",  Types.ANYTHING,
      "scale_factor", Types.FLOATING
  );

  @Getter
  private final List<Type> argumentTypes = getTypesFromArguments();

  @Getter
  private final Type returnType = Types.ANYTHING;

  @Override
  public ResultOrProblems<RiskscapeFunction> realize(RealizationContext context, FunctionCall functionCall,
      List<Type> argTypes) {
    Type toScaleType = argTypes.get(0);
    Type scaleFactorType = argTypes.get(1);

    if (! scaleFactorType.equals(Types.FLOATING)) {
      return ResultOrProblems.failed(ArgsProblems.mismatch(getArguments().get(1), scaleFactorType));
    }

    if (toScaleType.isNumeric()) {
      // Input is a simple numeric type
      return ResultOrProblems.of(new ScaleImpl(argTypes, Types.FLOATING, (x, scaleFactor) -> {
        return scaleNumber((Number)x, scaleFactor);
      }));
    }
    Optional<ScaleImpl> scaler = toScaleType.find(Struct.class)
        .flatMap(struct -> {
          Struct resultStruct = struct;
          List<StructMember> numericMembers = struct.getMembers().stream()
              .filter(member -> Nullable.strip(member.getType()).isNumeric())
              .collect(Collectors.toList());

          if (numericMembers.isEmpty()) {
            // we have to have some numeric members to scale
            return Optional.empty();
          }

          // update the struct type - scaling will change any integers to floating
          for (StructMember member : numericMembers) {
            boolean nullableMember = Nullable.is(member.getType());
            resultStruct = resultStruct.addOrReplace(member.getKey(),
                Nullable.ifTrue(nullableMember, Types.FLOATING));
          }
          Struct resultType = resultStruct;

          BiFunction<Object,Double,Object> structScaler = (x, scaleFactor) -> {
            Tuple input = (Tuple)x;
            Object[] argsArray = input.toArray();

            for (StructMember member : numericMembers) {
              argsArray[member.getIndex()] = scaleNumber((Number) argsArray[member.getIndex()], scaleFactor);
            }

            return Tuple.ofValues(resultType, argsArray);
          };
          return Optional.of(new ScaleImpl(argTypes, resultType, structScaler));
        });

    if (scaler.isPresent()) {
      return ResultOrProblems.of(scaler.get());
    }
    return ResultOrProblems.failed(ArgsProblems.get().notNumeric(getArguments().get(0), toScaleType));
  }

  private Double scaleNumber(Number toScale, double scaleFactor) {
    if (toScale == null) {
      return null;
    }
    return toScale.doubleValue() * scaleFactor;
  }

}
