/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.lucene.util;

import java.util.stream.IntStream;
import org.apache.lucene.internal.vectorization.VectorUtilSupport;
import org.apache.lucene.internal.vectorization.VectorizationProvider;

/**
 * Utilities for computations with numeric arrays, especially algebraic operations like vector dot
 * products. This class uses SIMD vectorization if the corresponding Java module is available and
 * enabled. To enable vectorized code, pass {@code --add-modules jdk.incubator.vector} to Java's
 * command line.
 *
 * <p>It will use CPU's <a href="https://en.wikipedia.org/wiki/Fused_multiply%E2%80%93add">FMA
 * instructions</a> if it is known to perform faster than separate multiply+add. This requires at
 * least Hotspot C2 enabled, which is the default for OpenJDK based JVMs.
 *
 * <p>To explicitly disable or enable FMA usage, pass the following system properties:
 *
 * <ul>
 *   <li>{@code -Dlucene.useScalarFMA=(auto|true|false)} for scalar operations
 *   <li>{@code -Dlucene.useVectorFMA=(auto|true|false)} for vectorized operations (with vector
 *       incubator module)
 * </ul>
 *
 * <p>The default is {@code auto}, which enables this for known CPU types and JVM settings. If
 * Hotspot C2 is disabled, FMA and vectorization are <strong>not</strong> used.
 *
 * <p>Vectorization and FMA is only supported for Hotspot-based JVMs; it won't work on OpenJ9-based
 * JVMs unless they provide {@link com.sun.management.HotSpotDiagnosticMXBean}. Please also make
 * sure that you have the {@code jdk.management} module enabled in modularized applications.
 */
public final class VectorUtil {

  private static final float EPSILON = 1e-4f;

  private static final VectorUtilSupport IMPL =
      VectorizationProvider.getInstance().getVectorUtilSupport();

  private VectorUtil() {}

  /**
   * Returns the vector dot product of the two vectors.
   *
   * @throws IllegalArgumentException if the vectors' dimensions differ.
   */
  public static float dotProduct(float[] a, float[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    float r = IMPL.dotProduct(a, b);
    assert Float.isFinite(r)
        : "not finite: "
            + r
            + " from <"
            + java.util.Arrays.toString(a)
            + ","
            + java.util.Arrays.toString(b)
            + ">";
    return r;
  }

  /**
   * Returns the cosine similarity between the two vectors.
   *
   * @throws IllegalArgumentException if the vectors' dimensions differ.
   */
  public static float cosine(float[] a, float[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    float r = IMPL.cosine(a, b);
    assert Float.isFinite(r);
    return r;
  }

  /** Returns the cosine similarity between the two vectors. */
  public static float cosine(byte[] a, byte[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    return IMPL.cosine(a, b);
  }

  /**
   * Returns the sum of squared differences of the two vectors.
   *
   * @throws IllegalArgumentException if the vectors' dimensions differ.
   */
  public static float squareDistance(float[] a, float[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    float r = IMPL.squareDistance(a, b);
    assert Float.isFinite(r);
    return r;
  }

  /** Returns the sum of squared differences of the two vectors. */
  public static int squareDistance(byte[] a, byte[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    return IMPL.squareDistance(a, b);
  }

  /** Returns the sum of squared differences between two uint4 (values between [0,15]) vectors. */
  public static int int4SquareDistance(byte[] a, byte[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    return IMPL.int4SquareDistance(a, b);
  }

  /**
   * Returns the sum of squared differences between two uint4 (values between [0,15]) vectors. The
   * second vector is considered "packed" (i.e. every byte representing two values).
   */
  public static int int4SquareDistanceSinglePacked(byte[] unpacked, byte[] packed) {
    if (packed.length != ((unpacked.length + 1) >> 1)) {
      throw new IllegalArgumentException(
          "vector dimensions differ: " + unpacked.length + "!= 2 * " + packed.length);
    }
    return IMPL.int4SquareDistanceSinglePacked(unpacked, packed);
  }

  /**
   * Returns the sum of squared differences between two uint4 (values between [0,15]) vectors. Both
   * vectors are considered "packed" (i.e. every byte representing two values).
   */
  public static int int4SquareDistanceBothPacked(byte[] a, byte[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    return IMPL.int4SquareDistanceBothPacked(a, b);
  }

  /** Returns the sum of squared differences of the two vectors where each byte is unsigned */
  public static int uint8SquareDistance(byte[] a, byte[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    return IMPL.uint8SquareDistance(a, b);
  }

  /**
   * Modifies the argument to be unit length, dividing by its l2-norm. IllegalArgumentException is
   * thrown for zero vectors.
   *
   * @return the input array after normalization
   */
  public static float[] l2normalize(float[] v) {
    l2normalize(v, true);
    return v;
  }

  public static boolean isUnitVector(float[] v) {
    double l1norm = IMPL.dotProduct(v, v);
    return Math.abs(l1norm - 1.0d) <= EPSILON;
  }

  /**
   * Modifies the argument to be unit length, dividing by its l2-norm.
   *
   * @param v the vector to normalize
   * @param throwOnZero whether to throw an exception when <code>v</code> has all zeros
   * @return the input array after normalization
   * @throws IllegalArgumentException when the vector is all zero and throwOnZero is true
   */
  public static float[] l2normalize(float[] v, boolean throwOnZero) {
    double l1norm = IMPL.dotProduct(v, v);
    if (l1norm == 0) {
      if (throwOnZero) {
        throw new IllegalArgumentException("Cannot normalize a zero-length vector");
      } else {
        return v;
      }
    }
    if (Math.abs(l1norm - 1.0d) <= EPSILON) {
      return v;
    }
    int dim = v.length;
    double l2norm = Math.sqrt(l1norm);
    for (int i = 0; i < dim; i++) {
      v[i] /= (float) l2norm;
    }
    return v;
  }

  /**
   * Adds the second argument to the first
   *
   * @param u the destination
   * @param v the vector to add to the destination
   */
  public static void add(float[] u, float[] v) {
    for (int i = 0; i < u.length; i++) {
      u[i] += v[i];
    }
  }

  /**
   * Dot product computed over signed bytes.
   *
   * @param a bytes containing a vector
   * @param b bytes containing another vector, of the same dimension
   * @return the value of the dot product of the two vectors
   */
  public static int dotProduct(byte[] a, byte[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    return IMPL.dotProduct(a, b);
  }

  /**
   * Dot product over bytes assuming that the values are actually unsigned.
   *
   * @param a uint8 byte vector
   * @param b another uint8 byte vector of the same dimension
   * @return the value of the dot product of the two vectors
   */
  public static int uint8DotProduct(byte[] a, byte[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    return IMPL.uint8DotProduct(a, b);
  }

  /**
   * Dot product computed over uint4 (values between [0,15]) bytes.
   *
   * @param a bytes containing a vector
   * @param b bytes containing another vector, of the same dimension
   * @return the value of the dot product of the two vectors
   */
  public static int int4DotProduct(byte[] a, byte[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    return IMPL.int4DotProduct(a, b);
  }

  /**
   * Dot product computed over uint4 (values between [0,15]) bytes. The second vector is considered
   * "packed" (i.e. every byte representing two values). The following packing is assumed:
   *
   * <pre class="prettyprint lang-java">
   *   packed[0] = (raw[0] * 16) | raw[packed.length];
   *   packed[1] = (raw[1] * 16) | raw[packed.length + 1];
   *   ...
   *   packed[packed.length - 1] = (raw[packed.length - 1] * 16) | raw[2 * packed.length - 1];
   * </pre>
   *
   * @param unpacked the unpacked vector, of even length
   * @param packed the packed vector, of length {@code (unpacked.length + 1) / 2}
   * @return the value of the dot product of the two vectors
   */
  public static int int4DotProductSinglePacked(byte[] unpacked, byte[] packed) {
    if (packed.length != ((unpacked.length + 1) >> 1)) {
      throw new IllegalArgumentException(
          "vector dimensions differ: " + unpacked.length + " != 2 * " + packed.length);
    }
    return IMPL.int4DotProductSinglePacked(unpacked, packed);
  }

  /**
   * Dot product computed over uint4 (values between [0,15]) bytes. Both vectors are considered
   * "packed" (i.e. every byte representing two values).
   *
   * @param a bytes containing a packed vector
   * @param b bytes containing another packed vector, of the same dimension
   * @return the value of the dot product of the two vectors
   */
  public static int int4DotProductBothPacked(byte[] a, byte[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException(
          "vector dimensions differ: " + a.length + " != " + b.length);
    }
    return IMPL.int4DotProductBothPacked(a, b);
  }

  /**
   * Dot product computed over int4 (values between [0,15]) bytes and a binary vector.
   *
   * @param q the int4 query vector
   * @param d the binary document vector
   * @return the dot product
   */
  public static long int4BitDotProduct(byte[] q, byte[] d) {
    if (q.length != d.length * 4) {
      throw new IllegalArgumentException(
          "vector dimensions incompatible: " + q.length + "!= " + 4 + " x " + d.length);
    }
    return IMPL.int4BitDotProduct(q, d);
  }

  /**
   * Dot product computed over int4 (values between [0,15]) bytes and a dibit (2-bit) vector.
   *
   * <p>The int4 query vector is expected to be transposed (4 stripes, each containing one bit plane
   * of all dimensions). The dibit document vector is expected to be transposed (2 stripes: lower
   * bits first, then upper bits).
   *
   * @param q the int4 query vector (transposed, 4 stripes)
   * @param d the dibit document vector (transposed, 2 stripes)
   * @return the dot product
   */
  public static long int4DibitDotProduct(byte[] q, byte[] d) {
    if (q.length != d.length * 2) {
      throw new IllegalArgumentException(
          "vector dimensions incompatible: " + q.length + "!= " + 2 + " x " + d.length);
    }
    return IMPL.int4DibitDotProduct(q, d);
  }

  /**
   * For xorBitCount we stride over the values as either 64-bits (long) or 32-bits (int) at a time.
   * On ARM Long::bitCount is not vectorized, and therefore produces less than optimal code, when
   * compared to Integer::bitCount. While Long::bitCount is optimal on x64. See
   * https://bugs.openjdk.org/browse/JDK-8336000
   */
  static final boolean XOR_BIT_COUNT_STRIDE_AS_INT = Constants.OS_ARCH.equals("aarch64");

  /**
   * XOR bit count computed over signed bytes.
   *
   * @param a bytes containing a vector
   * @param b bytes containing another vector, of the same dimension
   * @return the value of the XOR bit count of the two vectors
   */
  public static int xorBitCount(byte[] a, byte[] b) {
    if (a.length != b.length) {
      throw new IllegalArgumentException("vector dimensions differ: " + a.length + "!=" + b.length);
    }
    if (XOR_BIT_COUNT_STRIDE_AS_INT) {
      return xorBitCountInt(a, b);
    } else {
      return xorBitCountLong(a, b);
    }
  }

  /** XOR bit count striding over 4 bytes at a time. */
  static int xorBitCountInt(byte[] a, byte[] b) {
    int distance = 0, i = 0;
    for (final int upperBound = a.length & -Integer.BYTES; i < upperBound; i += Integer.BYTES) {
      distance +=
          Integer.bitCount(
              (int) BitUtil.VH_NATIVE_INT.get(a, i) ^ (int) BitUtil.VH_NATIVE_INT.get(b, i));
    }
    // tail:
    for (; i < a.length; i++) {
      distance += Integer.bitCount((a[i] ^ b[i]) & 0xFF);
    }
    return distance;
  }

  /** XOR bit count striding over 8 bytes at a time. */
  static int xorBitCountLong(byte[] a, byte[] b) {
    int distance = 0, i = 0;
    for (final int upperBound = a.length & -Long.BYTES; i < upperBound; i += Long.BYTES) {
      distance +=
          Long.bitCount(
              (long) BitUtil.VH_NATIVE_LONG.get(a, i) ^ (long) BitUtil.VH_NATIVE_LONG.get(b, i));
    }
    // tail:
    for (; i < a.length; i++) {
      distance += Integer.bitCount((a[i] ^ b[i]) & 0xFF);
    }
    return distance;
  }

  /**
   * Dot product score computed over signed bytes, scaled to be in [0, 1].
   *
   * @param a bytes containing a vector
   * @param b bytes containing another vector, of the same dimension
   * @return the value of the similarity function applied to the two vectors
   */
  public static float dotProductScore(byte[] a, byte[] b) {
    // divide by 2 * 2^14 (maximum absolute value of product of 2 signed bytes) * len
    float denom = (float) (a.length * (1 << 15));
    return 0.5f + dotProduct(a, b) / denom;
  }

  /**
   * @param vectorDotProductSimilarity the raw similarity between two vectors
   * @return A scaled score preventing negative scores for maximum-inner-product
   */
  public static float scaleMaxInnerProductScore(float vectorDotProductSimilarity) {
    if (vectorDotProductSimilarity < 0) {
      return 1 / (1 + -1 * vectorDotProductSimilarity);
    }
    return vectorDotProductSimilarity + 1;
  }

  /**
   * Converts a dot product or cosine similarity value to a normalized score in the [0, 1] range.
   *
   * <p>This transformation is necessary when consistent non-negative scores are required. It maps
   * input values from [-1, 1] to [0, 1] using the formula: (1 + value) / 2. Any result below 0 is
   * clamped to 0 for numerical safety.
   *
   * @param value the similarity value (dot product or cosine), typically in the range [-1, 1]
   * @return a normalized score between 0 and 1
   */
  public static float normalizeToUnitInterval(float value) {
    return Math.max((1 + value) / 2, 0);
  }

  /**
   * Maps a non-negative squared distance to a similarity score in the range (0, 1].
   *
   * <p>Uses the transformation: {@code similarity = 1 / (1 + squaredDistance)}. Smaller distances
   * yield scores closer to 1; larger distances approach 0.
   *
   * @param squaredDistance squared Euclidean distance (must be ≥ 0)
   * @return similarity score in (0, 1]
   */
  public static float normalizeDistanceToUnitInterval(float squaredDistance) {
    return 1.0f / (1.0f + squaredDistance);
  }

  /**
   * Checks if a float vector only has finite components.
   *
   * @param v bytes containing a vector
   * @return the vector for call-chaining
   * @throws IllegalArgumentException if any component of vector is not finite
   */
  public static float[] checkFinite(float[] v) {
    for (int i = 0; i < v.length; i++) {
      if (!Float.isFinite(v[i])) {
        throw new IllegalArgumentException("non-finite value at vector[" + i + "]=" + v[i]);
      }
    }
    return v;
  }

  /** Returns true if all dimensions of provided vector are zero, false otherwise. */
  public static boolean isZeroVector(float[] v) {
    for (float value : v) {
      if (value != 0) {
        return false;
      }
    }
    return true;
  }

  /** Returns true if all dimensions of provided vector are zero, false otherwise. */
  public static boolean isZeroVector(byte[] v) {
    for (byte value : v) {
      if (value != 0) {
        return false;
      }
    }
    return true;
  }

  /**
   * Given an array {@code buffer} that is sorted between indexes {@code 0} inclusive and {@code to}
   * exclusive, find the first array index whose value is greater than or equal to {@code target}.
   * This index is guaranteed to be at least {@code from}. If there is no such array index, {@code
   * to} is returned.
   */
  public static int findNextGEQ(int[] buffer, int target, int from, int to) {
    assert IntStream.range(0, to - 1).noneMatch(i -> buffer[i] > buffer[i + 1]);
    return IMPL.findNextGEQ(buffer, target, from, to);
  }

  /**
   * Scalar quantizes {@code vector}, putting the result into {@code dest}.
   *
   * @param vector the vector to quantize
   * @param dest the destination vector
   * @param scale the scaling factor
   * @param alpha the alpha value
   * @param minQuantile the lower quantile of the distribution
   * @param maxQuantile the upper quantile of the distribution
   * @return the corrective offset that needs to be applied to the score
   */
  public static float minMaxScalarQuantize(
      float[] vector, byte[] dest, float scale, float alpha, float minQuantile, float maxQuantile) {
    if (vector.length != dest.length)
      throw new IllegalArgumentException("source and destination arrays should be the same size");
    return IMPL.minMaxScalarQuantize(vector, dest, scale, alpha, minQuantile, maxQuantile);
  }

  /**
   * Recalculates the offset for {@code vector}.
   *
   * @param vector the vector to quantize
   * @param oldAlpha the previous alpha value
   * @param oldMinQuantile the previous lower quantile
   * @param scale the scaling factor
   * @param alpha the alpha value
   * @param minQuantile the lower quantile of the distribution
   * @param maxQuantile the upper quantile of the distribution
   * @return the new corrective offset
   */
  public static float recalculateOffset(
      byte[] vector,
      float oldAlpha,
      float oldMinQuantile,
      float scale,
      float alpha,
      float minQuantile,
      float maxQuantile) {
    return IMPL.recalculateScalarQuantizationOffset(
        vector, oldAlpha, oldMinQuantile, scale, alpha, minQuantile, maxQuantile);
  }

  /**
   * filter both {@code docBuffer} and {@code scoreBuffer} with {@code minScoreInclusive}, each
   * {@code docBuffer} and {@code scoreBuffer} of the same index forms a pair, pairs with score not
   * greater than or equal to {@code minScoreInclusive} will be filtered out from the array.
   *
   * @param docBuffer doc buffer contains docs (or some other value forms a pair with {@code
   *     scoreBuffer})
   * @param scoreBuffer score buffer contains scores to be compared with {@code minScoreInclusive}
   * @param minScoreInclusive minimal required score to not be filtered out
   * @param upTo where the filter should end
   * @return how many pairs left after filter
   */
  public static int filterByScore(
      int[] docBuffer, double[] scoreBuffer, double minScoreInclusive, int upTo) {
    if (docBuffer.length != scoreBuffer.length || docBuffer.length < upTo) {
      throw new IllegalArgumentException(
          "docBuffer and scoreBuffer should keep same length and at least as long as upTo");
    }
    return IMPL.filterByScore(docBuffer, scoreBuffer, minScoreInclusive, upTo);
  }
}
