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

import static java.lang.foreign.ValueLayout.JAVA_BYTE;
import static java.lang.foreign.ValueLayout.JAVA_DOUBLE_UNALIGNED;
import static org.apache.datasketches.common.Util.LS;
import static org.apache.datasketches.sampling.PreambleUtil.EBPPS_SER_VER;
import static org.apache.datasketches.sampling.PreambleUtil.EMPTY_FLAG_MASK;
import static org.apache.datasketches.sampling.PreambleUtil.HAS_PARTIAL_ITEM_MASK;

import java.lang.foreign.MemorySegment;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import org.apache.datasketches.common.ArrayOfItemsSerDe;
import org.apache.datasketches.common.Family;
import org.apache.datasketches.common.SketchesArgumentException;

/**
 * An implementation of an Exact and Bounded Sampling Proportional to Size sketch.
 *
 * <p>From: "Exact PPS Sampling with Bounded Sample Size",
 * B. Hentschel, P. J. Haas, Y. Tian. Information Processing Letters, 2023.
 *
 * <p>This sketch samples data from a stream of items proportional to the weight of each item.
 * The sample guarantees the presence of an item in the result is proportional to that item's
 * portion of the total weight seen by the sketch, and returns a sample no larger than size k.
 *
 * <p>The sample may be smaller than k and the resulting size of the sample potentially includes
 * a probabilistic component, meaning the resulting sample size is not always constant.
 * @param <T> the item class type
 * @author Jon Malkin
 */
public final class EbppsItemsSketch<T> {
  private static final int MAX_K = Integer.MAX_VALUE - 2;
  private static final int EBPPS_C_DOUBLE        = 40; // part of sample state, not preamble
  private static final int EBPPS_ITEMS_START     = 48;

  private int k_;                      // max size of sketch, in items
  private long n_;                     // total number of items processed by the sketch

  private double cumulativeWt_;        // total weight of items processed by the sketch
  private double wtMax_;               // maximum weight seen so far
  private double rho_;                 // latest scaling parameter for downsampling

  private EbppsItemsSample<T> sample_; // Object holding the current state of the sample

  final private EbppsItemsSample<T> tmp_;    // temporary storage

  /**
   * Constructor
   * @param k The maximum number of samples to retain
   */
  public EbppsItemsSketch(final int k) {
    checkK(k);
    k_ = k;
    rho_ = 1.0;
    sample_ = new EbppsItemsSample<>(k);
    tmp_ = new EbppsItemsSample<>(1);
  }

  // private copy constructor
  private EbppsItemsSketch(final EbppsItemsSketch<T> other) {
    k_ = other.k_;
    n_ = other.n_;
    rho_ = other.rho_;
    cumulativeWt_ = other.cumulativeWt_;
    wtMax_ = other.wtMax_;
    sample_ = new EbppsItemsSample<>(other.sample_);
    tmp_ = new EbppsItemsSample<>(1);
  }

  // private constructor for heapify
  private EbppsItemsSketch(final EbppsItemsSample<T> sample,
                           final int k,
                           final long n,
                           final double cumWt,
                           final double maxWt,
                           final double rho) {
    k_ = k;
    n_ = n;
    cumulativeWt_ = cumWt;
    wtMax_ = maxWt;
    rho_ = rho;
    sample_ = sample;
    tmp_ = new EbppsItemsSample<>(1);
  }

  /**
   * Returns a sketch instance of this class from the given srcSeg,
   * which must be a MemorySegment representation of this sketch class.
   *
   * @param <T>    The type of item this sketch contains
   * @param srcSeg a MemorySegment representation of a sketch of this class.
   * @param serDe  An instance of ArrayOfItemsSerDe
   * @return a sketch instance of this class
   */
  public static <T> EbppsItemsSketch<T> heapify(final MemorySegment srcSeg, final ArrayOfItemsSerDe<T> serDe)
  {
    final int numPreLongs = PreambleUtil.getAndCheckPreLongs(srcSeg);
    final int serVer = PreambleUtil.extractSerVer(srcSeg);
    final int familyId = PreambleUtil.extractFamilyID(srcSeg);
    final int flags = PreambleUtil.extractFlags(srcSeg);
    final boolean isEmpty = (flags & EMPTY_FLAG_MASK) != 0;
    final boolean hasPartialItem = (flags & HAS_PARTIAL_ITEM_MASK) != 0;

    // Check values
    if (isEmpty) {
      if (numPreLongs != Family.EBPPS.getMinPreLongs()) {
        throw new SketchesArgumentException("Possible corruption: Must be " + Family.EBPPS.getMinPreLongs()
                + " for an empty sketch. Found: " + numPreLongs);
      }
    } else if (numPreLongs != Family.EBPPS.getMaxPreLongs()) {
      throw new SketchesArgumentException("Possible corruption: Must be "
              + Family.EBPPS.getMaxPreLongs() + " for a non-empty sketch. Found: " + numPreLongs);
    }
    if (serVer != EBPPS_SER_VER) {
        throw new SketchesArgumentException(
                "Possible Corruption: Ser Ver must be " + EBPPS_SER_VER + ": " + serVer);
    }
    final int reqFamilyId = Family.EBPPS.getID();
    if (familyId != reqFamilyId) {
      throw new SketchesArgumentException(
              "Possible Corruption: FamilyID must be " + reqFamilyId + ": " + familyId);
    }

    final int k = PreambleUtil.extractK(srcSeg);
    if ((k < 1) || (k > MAX_K)) {
      throw new SketchesArgumentException("Possible Corruption: k must be at least 1 "
              + "and less than " + MAX_K + ". Found: " + k);
    }

    if (isEmpty) {
      return new EbppsItemsSketch<>(k);
    }

    final long n = PreambleUtil.extractN(srcSeg);
    if (n < 0) {
      throw new SketchesArgumentException("Possible Corruption: n cannot be negative: " + n);
    }

    final double cumWt = PreambleUtil.extractEbppsCumulativeWeight(srcSeg);
    if ((cumWt < 0.0) || Double.isNaN(cumWt) || Double.isInfinite(cumWt)) {
      throw new SketchesArgumentException("Possible Corruption: cumWt must be nonnegative and finite: " + cumWt);
    }

    final double maxWt = PreambleUtil.extractEbppsMaxWeight(srcSeg);
    if ((maxWt < 0.0) || Double.isNaN(maxWt) || Double.isInfinite(maxWt)) {
      throw new SketchesArgumentException("Possible Corruption: maxWt must be nonnegative and finite: " + maxWt);
    }

    final double rho = PreambleUtil.extractEbppsRho(srcSeg);
    if ((rho < 0.0) || (rho > 1.0) ||  Double.isNaN(rho) || Double.isInfinite(rho)) {
      throw new SketchesArgumentException("Possible Corruption: rho must be in [0.0, 1.0]: " + rho);
    }

    // extract C (part of sample_, not the preamble)
    // due to numeric precision issues, c may occasionally be very slightly larger than k
    final double c = srcSeg.get(JAVA_DOUBLE_UNALIGNED, EBPPS_C_DOUBLE);
    if ((c < 0) || (c >= (k + 1)) || Double.isNaN(c) || Double.isInfinite(c)) {
      throw new SketchesArgumentException("Possible Corruption: c must be between 0 and k: " + c);
    }

    // extract items
    final int numTotalItems = (int) Math.ceil(c);
    final int numFullItems = (int) Math.floor(c); // floor() not strictly necessary
    final int offsetBytes = EBPPS_ITEMS_START;
    final T[] rawItems = serDe.deserializeFromMemorySegment(srcSeg.asSlice(offsetBytes), 0, numTotalItems);
    final List<T> itemsList = Arrays.asList(rawItems);
    final ArrayList<T> data;
    final T partialItem;
    if (hasPartialItem) {
      if (numFullItems >= numTotalItems) {
        throw new SketchesArgumentException("Possible Corruption: Expected partial item but none found");
      }

      data = new ArrayList<>(itemsList.subList(0, numFullItems));
      partialItem = itemsList.get(numFullItems); // 0-based, so last item
    } else {
      data = new ArrayList<>(itemsList);
      partialItem = null; // just to be explicit
    }

    final EbppsItemsSample<T> sample = new EbppsItemsSample<>(data, partialItem, c);

    return new EbppsItemsSketch<>(sample, k, n, cumWt, maxWt, rho);
  }

  /**
   * Updates this sketch with the given data item with weight 1.0.
   * @param item an item from a stream of items
   */
  public void update(final T item) {
    update(item, 1.0);
  }

  /**
   * Updates this sketch with the given data item with the given weight.
   * @param item an item from a stream of items
   * @param weight the weight of the item
   */
  public void update(final T item, final double weight) {
    if ((weight < 0.0) || Double.isNaN(weight) || Double.isInfinite(weight)) {
      throw new SketchesArgumentException("Item weights must be nonnegative and finite. "
        + "Found: " + weight);
    }
    if (weight == 0.0) {
      return;
    }

    final double newCumWt = cumulativeWt_ + weight;
    final double newWtMax = Math.max(wtMax_, weight);
    final double newRho = Math.min(1.0 / newWtMax, k_ / newCumWt);

    if (cumulativeWt_ > 0.0) {
      sample_.downsample((newRho / rho_));
    }

    tmp_.replaceContent(item, newRho * weight);
    sample_.merge(tmp_);

    cumulativeWt_ = newCumWt;
    wtMax_ = newWtMax;
    rho_ = newRho;
    ++n_;
  }

  /* Merging
   * There is a trivial merge algorithm that involves downsampling each sketch A and B
   * as A.cum_wt / (A.cum_wt + B.cum_wt) and B.cum_wt / (A.cum_wt + B.cum_wt),
   * respectively. That merge does preserve first-order probabilities, specifically
   * the probability proportional to size property, and like all other known merge
   * algorithms distorts second-order probabilities (co-occurrences). There are
   * pathological cases, most obvious with k=2 and A.cum_wt == B.cum_wt where that
   * approach will always take exactly 1 item from A and 1 from B, meaning the
   * co-occurrence rate for two items from either sketch is guaranteed to be 0.0.
   *
   * With EBPPS, once an item is accepted into the sketch we no longer need to
   * track the item's weight: All accepted items are treated equally. As a result, we
   * can take inspiration from the reservoir sampling merge in the datasketches-java
   * library. We need to merge the smaller sketch into the larger one, swapping as
   * needed to ensure that, at which point we simply call update() with the items
   * in the smaller sketch as long as we adjust the weight appropriately.
   * Merging smaller into larger is essential to ensure that no item has a
   * contribution to C > 1.0.
   */

  /**
   * Merges the provided sketch into the current one.
   * @param other the sketch to merge into the current object
   */
  public void merge(final EbppsItemsSketch<T> other) {
    if (other.getCumulativeWeight() == 0.0) { /* do nothing */ }
    else if (other.getCumulativeWeight() > cumulativeWt_) {
      // need to swap this with other
      // make a copy of other, merge into it, and take the result
      final EbppsItemsSketch<T> copy = new EbppsItemsSketch<>(other);
      copy.internalMerge(this);
      k_ = copy.k_;
      n_ = copy.n_;
      cumulativeWt_ = copy.cumulativeWt_;
      wtMax_ = copy.wtMax_;
      rho_ = copy.rho_;
      sample_ = copy.sample_;
    } else {
      internalMerge(other);
    }
  }

  // merge implementation called exclusively from public merge()
  private void internalMerge(final EbppsItemsSketch<T> other) {
    // assumes that other.cumulativeWeight_ <= cumulativeWt_m
    // which must be checked before calling this

    final double finalCumWt = cumulativeWt_ + other.cumulativeWt_;
    final double newWtMax = Math.max(wtMax_, other.wtMax_);
    k_ = Math.min(k_, other.k_);
    final long newN = n_ + other.n_;

    // Insert other's items with the cumulative weight
    // split between the input items. We repeat the same process
    // for full items and the partial item, scaling the input
    // weight appropriately.
    // We handle all C input items, meaning we always process
    // the partial item using a scaled down weight.
    // Handling the partial item by probabilistically including
    // it as a full item would be correct on average but would
    // introduce bias for any specific merge operation.
    final double avgWt = other.cumulativeWt_ / other.getC();
    final ArrayList<T> items = other.sample_.getFullItems();
    if (items != null) {
      for (final T item : items) {
        // newWtMax is pre-computed
        final double newCumWt = cumulativeWt_ + avgWt;
        final double newRho = Math.min(1.0 / newWtMax, k_ / newCumWt);

        if (cumulativeWt_ > 0.0) {
          sample_.downsample(newRho / rho_);
        }

        tmp_.replaceContent(item, newRho * avgWt);
        sample_.merge(tmp_);

        cumulativeWt_ = newCumWt;
        rho_ = newRho;
      }
    }

    // insert partial item with weight scaled by the fractional part of C
    if (other.sample_.hasPartialItem()) {
      final double otherCFrac = other.getC() % 1;
      final double newCumWt = cumulativeWt_ + (otherCFrac * avgWt);
      final double newRho = Math.min(1.0 / newWtMax, k_ / newCumWt);

      if (cumulativeWt_ > 0.0) {
        sample_.downsample(newRho / rho_);
      }

      tmp_.replaceContent(other.sample_.getPartialItem(), newRho * otherCFrac * avgWt);
      sample_.merge(tmp_);

      // cumulativeWt_ will be assigned momentarily
      rho_ = newRho;
    }

    // avoid numeric issues by setting cumulative weight to the
    // pre-computed value
    cumulativeWt_ = finalCumWt;
    n_ = newN;
  }

  /**
   * Returns a copy of the current sample. The exact size may be
   * probabilistic, differing by at most 1 item.
   * @return the current sketch sample
   */
  public ArrayList<T> getResult() { return sample_.getSample(); }

  /**
   * Provides a human-readable summary of the sketch
   * @return a summary of information in the sketch
   */
  @Override
   public String toString() {
    final StringBuilder sb = new StringBuilder();

    sb.append(LS);
    final String thisSimpleName = this.getClass().getSimpleName();
    sb.append("### ").append(thisSimpleName).append(" SUMMARY: ").append(LS);
    sb.append("   k            : ").append(k_).append(LS);
    sb.append("   n            : ").append(n_).append(LS);
    sb.append("   Cum. weight  : ").append(cumulativeWt_).append(LS);
    sb.append("   wtMax        : ").append(wtMax_).append(LS);
    sb.append("   rho          : ").append(rho_).append(LS);
    sb.append("   C            : ").append(sample_.getC()).append(LS);
    sb.append("### END SKETCH SUMMARY").append(LS);

    return sb.toString();
  }

  /**
   * Returns the configured maximum sample size.
   * @return configured maximum sample size
   */
  public int getK() { return k_; }

  /**
   * Returns the number of items processed by the sketch, regardless
   * of item weight.
   * @return count of items processed by the sketch
   */
  public long getN() { return n_; }

  /**
   * Returns the cumulative weight of items processed by the sketch.
   * @return cumulative weight of items seen
   */
  public double getCumulativeWeight() { return cumulativeWt_; }

  /**
   * Returns the expected number of samples returned upon a call to
   * getResult(). The number is a floating point value, where the
   * fractional portion represents the probability of including a
   * "partial item" from the sample.
   *
   * <p>The value C should be no larger than the sketch's configured
   * value of k, although numerical precision limitations mean it
   * may exceed k by double precision floating point error margins
   * in certain cases.
   * @return The expected number of samples returned when querying the sketch
   */
  public double getC() { return sample_.getC(); }

  /**
   * Returns true if the sketch is empty.
   * @return empty flag
   */
  public boolean isEmpty() { return n_ == 0; }

  /**
   * Resets the sketch to its default, empty state.
   */
  public void reset() {
    n_ = 0;
    cumulativeWt_ = 0.0;
    wtMax_ = 0.0;
    rho_ = 1.0;
    sample_ = new EbppsItemsSample<>(k_);
  }

  /**
   * Returns the size of a byte array representation of this sketch. May fail for polymorphic item types.
   *
   * @param serDe An instance of ArrayOfItemsSerDe
   * @return the length of a byte array representation of this sketch
   */
  public int getSerializedSizeBytes(final ArrayOfItemsSerDe<? super T> serDe) {
    if (isEmpty()) {
      return Family.EBPPS.getMinPreLongs() << 3;
    } else if (sample_.getC() < 1.0) {
      return getSerializedSizeBytes(serDe, sample_.getPartialItem().getClass());
    } else {
      return getSerializedSizeBytes(serDe, sample_.getSample().get(0).getClass());
    }
  }

  /**
   * Returns the length of a byte array representation of this sketch. Copies contents into an array of the
   * specified class for serialization to allow for polymorphic types.
   *
   * @param serDe An instance of ArrayOfItemsSerDe
   * @param clazz The class represented by &lt;T&gt;
   * @return the length of a byte array representation of this sketch
   */
  public int getSerializedSizeBytes(final ArrayOfItemsSerDe<? super T> serDe, final Class<?> clazz) {
    if (n_ == 0) {
      return Family.EBPPS.getMinPreLongs() << 3;
    }

    final int preLongs = Family.EBPPS.getMaxPreLongs();
    final byte[] itemBytes = serDe.serializeToByteArray(sample_.getAllSamples(clazz));
    // in C++, c_ is serialized as part of the sample_ and not included in the header size
    return (preLongs << 3) + Double.BYTES + itemBytes.length;
  }

  /**
   * Returns a byte array representation of this sketch. May fail for polymorphic item types.
   *
   * @param serDe An instance of ArrayOfItemsSerDe
   * @return a byte array representation of this sketch
   */
  public byte[] toByteArray(final ArrayOfItemsSerDe<? super T> serDe) {
    if (n_ == 0) {
      // null class is ok since empty -- no need to call serDe
      return toByteArray(serDe, null);
    } else if (sample_.getC() < 1.0) {
      return toByteArray(serDe, sample_.getPartialItem().getClass());
    } else {
      return toByteArray(serDe, sample_.getSample().get(0).getClass());
    }
  }

   /**
   * Returns a byte array representation of this sketch. Copies contents into an array of the
   * specified class for serialization to allow for polymorphic types.
   *
   * @param serDe An instance of ArrayOfItemsSerDe
   * @param clazz The class represented by &lt;T&gt;
   * @return a byte array representation of this sketch
   */
  public byte[] toByteArray(final ArrayOfItemsSerDe<? super T> serDe, final Class<?> clazz) {
    final int preLongs, outBytes;
    final boolean empty = n_ == 0;
    byte[] itemBytes = null; // for serialized items from sample_

    if (empty) {
      preLongs = 1;
      outBytes = 8;
    } else {
      preLongs = Family.EBPPS.getMaxPreLongs();
      itemBytes = serDe.serializeToByteArray(sample_.getAllSamples(clazz));
      // in C++, c_ is serialized as part of the sample_ and not included in the header size
      outBytes = (preLongs << 3) + Double.BYTES + itemBytes.length;
    }
    final byte[] outArr = new byte[outBytes];
    final MemorySegment seg = MemorySegment.ofArray(outArr);

    // Common header elements
    PreambleUtil.insertPreLongs(seg, preLongs);              // Byte 0
    PreambleUtil.insertSerVer(seg, EBPPS_SER_VER);           // Byte 1
    PreambleUtil.insertFamilyID(seg, Family.EBPPS.getID());  // Byte 2
    if (empty) {
      PreambleUtil.insertFlags(seg, EMPTY_FLAG_MASK);        // Byte 3
    } else {
      PreambleUtil.insertFlags(seg, sample_.hasPartialItem() ? HAS_PARTIAL_ITEM_MASK : 0);
    }
    PreambleUtil.insertK(seg, k_);                           // Bytes 4-7

    // conditional elements
    if (!empty) {
      PreambleUtil.insertN(seg, n_);
      PreambleUtil.insertEbppsCumulativeWeight(seg, cumulativeWt_);
      PreambleUtil.insertEbppsMaxWeight(seg, wtMax_);
      PreambleUtil.insertEbppsRho(seg, rho_);

      // data from sample_ -- itemBytes includes the partial item
      seg.set(JAVA_DOUBLE_UNALIGNED, EBPPS_C_DOUBLE, sample_.getC());
      MemorySegment.copy(itemBytes, 0, seg, JAVA_BYTE, EBPPS_ITEMS_START, itemBytes.length);
    }

    return outArr;
  }

  private static void checkK(final int k) {
    if ((k <= 0) || (k > MAX_K)) {
      throw new SketchesArgumentException("k must be strictly positive and less than " + MAX_K);
    }
  }
}
