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


import java.lang.reflect.Array;
import java.util.Arrays;

import javax.annotation.Nullable;

import nz.org.riskscape.engine.Tuple;
import nz.org.riskscape.engine.rl.RealizedExpression;
import nz.org.riskscape.engine.rl.ScopedLambdaExpression;
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;

/**
 * Base class for implementing stackable continuous functions.
 */
abstract class StackableContinuousFunctionType<F extends StackableContinuousState> extends ContinuousFunctionType {

  /**
   * If true, the numeric tuples this function creates are stored as floats, packed in to a float[] to reduce memory
   * usage
   */
  private final boolean compress;

  /**
   * The number of members in the tuples this function creates
   */
  private final int numCompressedMembers;

  private final Struct uncompressedType;

  StackableContinuousFunctionType(int dimensions, RealizedExpression valueExpression, Type returnType,
      boolean compress) {
    super(dimensions, valueExpression, returnType);

    /*
     * Sanity check the return type that it's either a struct of floating or an individual floating value.
     * Note that we ignore the interpolant expression's return type and use the function's return type - this
     * should be fine as they will all end up being floating point numbers anyway.
     */

    this.compress = compress;
    this.uncompressedType = returnType.find(Struct.class).orElse(null);
    if (uncompressedType == null) {
      if (returnType.equals(Types.FLOATING)) {
        this.numCompressedMembers = 1;
      } else {
        this.numCompressedMembers = 0;
      }
    } else if (compress) {
      for (StructMember member : uncompressedType.getMembers()) {
        if (member.getType() != Types.FLOATING) {
          throw new IllegalArgumentException("Result type not suitable for compression: " + returnType);
        }
      }

      this.numCompressedMembers = uncompressedType.size();
    } else {
      numCompressedMembers = 0;
    }

    if (compress && numCompressedMembers == 0) {
      throw new IllegalArgumentException("Result type not suitable for compression: " + returnType);
    }
  }

  /**
   * Get the value from function for the given key.
   *
   * If function holds an already computed value for the key then that should be returned. Otherwise
   * the value should be computed and saved in function of reuse.
   */
  abstract Object getOrComputeValue(F function, int key);

  /**
   * Pre-emptively populate the given function with values - needs to be done before a function can be 'stacked'
   */
  void populate(F function) {
    // we stack from highest to lowest as this gives any zero loss short circuits the best chance
    // of improving performance.
    for (int i = getSize() - 1; i >= 0; i--) {
      getOrComputeValue(function, i);
    }

    // can now garbage collect this - it's potentially a largish amount of data and if we're doing an stack_continuous
    // aggregation these can begin to add up during the accumulation part.
    function.lambda = null;
  }

  /**
   * @return a new instance of a {@link StackableContinuousState} that closes over the given scope. The resulting
   * function can be passed around in a pipeline to compute interpolated values at a later time.
   */
  abstract F newFunction(@Nullable ScopedLambdaExpression scope);

  /**
   * @return a new instance of a {@link StackableContinuousState} that can be used for accumulating values (but
   * can not be used for computing new values - it does not have access to any expression scope)
   */
  F newAccumulator() {
    return newFunction(null);
  }

  /*
   * Get/set helpers that will eventually handle deserialization?
   */

  /**
   * @return a new backing array for storing values of the type this function produces.  This value gets assigned to the
   * `values` field on StackableContinuousState and is the array that `setValue` and `getValue` ultimately operate on
   */
  public Object newArray() {
    final int size = getSize();
    if (compress) {
      float[] floats = new float[size * numCompressedMembers];

      // The NaN value is used as a substitute for null.  If the interpolant function ends up producing nan values, then
      // we will end up recomputing the loss needlessly
      Arrays.fill(floats, Float.NaN);
      return floats;
    } else {
      return new Object[size];
    }
  }

  /**
   * Set a computed value for the given key.
   */
  public void setValue(F function, int arrayIndex, Object value) {

    if (compress) {
      // serialize our tuple of (hopefully) numbers in to floats
      float[] floats = (float[]) function.values;
      if (uncompressedType != null) {
        Tuple tuple = (Tuple) value;

        // offset in to the floats array
        final int offset = arrayIndex * numCompressedMembers;

        // read through numCompressedMembers from the tuple and store them contiguously
        for (int memberIndex = 0; memberIndex < numCompressedMembers; memberIndex++) {
          Number number = tuple.fetch(memberIndex);
          floats[offset + memberIndex] = number.floatValue();
        }
      } else {
        floats[arrayIndex] = ((Double) value).floatValue();
      }
    } else {
      ((Object[]) function.values)[arrayIndex] = value;
    }
  }

  /**
   * Get the previously computed value for the given key.
   */
  public Object getValue(F function, int index) {
    if (compress) {
      // deserialize floats back in to a tuple
      float[] floats = (float[]) function.values;

      if (uncompressedType != null) {
        // struct result case
        Tuple tuple = Tuple.of(uncompressedType);

        // offset in to the floats array
        int offset = index * numCompressedMembers;

        // NB: we are using nan as a stand in for null - if the interpolant function we're given returns NaNs,
        // performance
        // is gonna suck
        if (Float.isNaN(floats[offset])) {
          return null;
        }

        for (StructMember member : uncompressedType.getMembers()) {
          float f = floats[offset++];
          // we assert that the result type only ever contains doubles in the constructor, as that's the only thing
          // (I think) a continuous function will ever return (Because of interpolation)
          tuple.set(member, Double.valueOf(f));
        }

        return tuple;
      } else {
        float value = floats[index];
        // nan is our null stand in
        if (Float.isNaN(value)) {
          return null;
        } else {
          return Double.valueOf(value);
        }
      }
    } else {
      return ((Object[]) function.values)[index];
    }
  }

  /**
   * @return the number of unique values this function can store in it.  This corresponds to the conceptual length of
   * the array - `index` values passed to the `getValue` and `setValue` methods must be less that the size's return
   * value.
   */
  public abstract int getSize();

  /**
   * Set all the values from other into this function replacing any values that exist already.
   */
  public void copyValues(F srcFunction, F destFunction) {
    System.arraycopy(srcFunction.values, 0, destFunction.values, 0, Array.getLength(destFunction.values));
  }

}
