package io.prometheus.metrics.core.metrics;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
import static org.assertj.core.data.Offset.offset;

import io.prometheus.metrics.core.metrics.CKMSQuantiles.Quantile;
import java.util.*;
import org.apache.commons.math3.distribution.NormalDistribution;
import org.apache.commons.math3.random.JDKRandomGenerator;
import org.apache.commons.math3.random.RandomGenerator;
import org.junit.jupiter.api.Test;

class CKMSQuantilesTest {

  private final Quantile qMin = new Quantile(0.0, 0.00);
  private final Quantile q50 = new Quantile(0.5, 0.01);
  private final Quantile q95 = new Quantile(0.95, 0.005);
  private final Quantile q99 = new Quantile(0.99, 0.001);
  private final Quantile qMax = new Quantile(1.0, 0.00);

  @Test
  public void testGetOnEmptyValues() {
    CKMSQuantiles ckms = new CKMSQuantiles(q50, q95, q99);
    assertThat(Double.isNaN(ckms.get(q95.quantile))).isTrue();
  }

  @Test
  public void testGet() {
    Random random = new Random(0);
    CKMSQuantiles ckms = new CKMSQuantiles(q50, q95, q99);
    List<Double> input = shuffledValues(100, random);
    for (double value : input) {
      ckms.insert(value);
    }
    validateResults(ckms);
  }

  @Test
  public void testBatchInsert() {
    Random random = new Random(1);
    testInsertBatch(1, 1, 100, random);
    testInsertBatch(1, 10, 100, random);
    testInsertBatch(2, 10, 100, random);
    testInsertBatch(
        2, 110, 100,
        random); // compress never called, because compress interval > number of inserts
    testInsertBatch(3, 10, 100, random);
    testInsertBatch(10, 10, 100, random);
    testInsertBatch(128, 128, 1, random);
    testInsertBatch(128, 128, 1000, random);
    testInsertBatch(128, 128, 10 * 1000, random);
    testInsertBatch(128, 128, 100 * 1000, random);
    testInsertBatch(128, 128, 1000 * 1000, random);
  }

  private void testInsertBatch(
      int batchSize, int compressInterval, int totalNumber, Random random) {
    System.out.println(
        "testInsertBatch(batchSize="
            + batchSize
            + ", compressInterval="
            + compressInterval
            + ", totalNumber="
            + totalNumber
            + ")");
    CKMSQuantiles ckms = new CKMSQuantiles(q50, q95);
    int insertsSinceCompress = 0;
    List<Double> input = shuffledValues(totalNumber, random);
    for (int i = 0; i < input.size(); i += batchSize) {
      double[] batch = new double[batchSize];
      int j;
      for (j = 0; j < batchSize && i + j < input.size(); j++) {
        batch[j] = input.get(i + j);
      }
      Arrays.sort(batch, 0, j);
      ckms.insertBatch(batch, j);
      validateSamples(ckms); // after each insert the samples should still be valid
      insertsSinceCompress += j;
      if (insertsSinceCompress >= compressInterval) {
        ckms.compress();
        validateSamples(ckms); // after each compress the samples should still be valid
        insertsSinceCompress = 0;
      }
    }
    validateResults(ckms);
  }

  @Test
  public void testGetWithAMillionElements() {
    Random random = new Random(2);
    List<Double> input = shuffledValues(1000 * 1000, random);
    CKMSQuantiles ckms = new CKMSQuantiles(q50, q95, q99);
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
    assertThat(ckms.samples).as("sample size should be way below 1_000_000").hasSizeLessThan(1000);
  }

  @Test
  public void testMin() {
    Random random = new Random(3);
    List<Double> input = shuffledValues(1000, random);
    CKMSQuantiles ckms = new CKMSQuantiles(qMin);
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
    ckms.compress();
    assertThat(ckms.samples).hasSize(2);
  }

  @Test
  public void testMax() {
    Random random = new Random(4);
    List<Double> input = shuffledValues(1000, random);
    CKMSQuantiles ckms = new CKMSQuantiles(qMax);
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
    ckms.compress();
    assertThat(ckms.samples).hasSize(2);
  }

  @Test
  public void testMinMax() {
    Random random = new Random(5);
    List<Double> input = shuffledValues(1000, random);
    CKMSQuantiles ckms = new CKMSQuantiles(qMin, qMax);
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
    ckms.compress();
    assertThat(ckms.samples).hasSize(2);
  }

  @Test
  public void testMinAndOthers() {
    Random random = new Random(6);
    List<Double> input = shuffledValues(1000, random);
    CKMSQuantiles ckms = new CKMSQuantiles(q95, qMin);
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
    assertThat(ckms.samples).hasSizeLessThan(200); // should be a lot less than input.size()
  }

  @Test
  public void testMaxAndOthers() {
    Random random = new Random(7);
    List<Double> input = shuffledValues(10000, random);
    CKMSQuantiles ckms = new CKMSQuantiles(q50, q95, qMax);
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
    assertThat(ckms.samples).hasSizeLessThan(200); // should be a lot less than input.size()
  }

  @Test
  public void testMinMaxAndOthers() {
    Random random = new Random(8);
    List<Double> input = shuffledValues(10000, random);
    CKMSQuantiles ckms = new CKMSQuantiles(qMin, q50, q95, q99, qMax);
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
    assertThat(ckms.samples).hasSizeLessThan(200); // should be a lot less than input.size()
  }

  @Test
  public void testExactQuantile() {
    Random random = new Random(9);
    List<Double> input = shuffledValues(10000, random);
    CKMSQuantiles ckms = new CKMSQuantiles(new Quantile(0.95, 0));
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
    // With epsilon == 0 we need to keep all inputs in samples.
    assertThat(ckms.samples).hasSameSizeAs(input);
  }

  @Test
  public void testExactAndOthers() {
    Random random = new Random(10);
    List<Double> input = shuffledValues(10000, random);
    CKMSQuantiles ckms = new CKMSQuantiles(q50, new Quantile(0.95, 0), q99);
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
    // With epsilon == 0 we need to keep all inputs in samples.
    assertThat(ckms.samples).hasSameSizeAs(input);
  }

  @Test
  public void testExactAndMin() {
    Random random = new Random(11);
    List<Double> input = shuffledValues(10000, random);
    CKMSQuantiles ckms = new CKMSQuantiles(qMin, q50, new Quantile(0.95, 0));
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
    // With epsilon == 0 we need to keep all inputs in samples.
    assertThat(ckms.samples).hasSameSizeAs(input);
  }

  @Test
  public void testMaxEpsilon() {
    Random random = new Random(12);
    List<Double> input = shuffledValues(10000, random);
    // epsilon == 1 basically gives you random results, but it should still not throw an exception.
    CKMSQuantiles ckms = new CKMSQuantiles(new Quantile(0.95, 1));
    for (double v : input) {
      ckms.insert(v);
    }
    validateResults(ckms);
  }

  @Test
  public void testGetGaussian() {
    RandomGenerator rand = new JDKRandomGenerator();
    rand.setSeed(0);

    double mean = 0.0;
    double stddev = 1.0;
    NormalDistribution normalDistribution =
        new NormalDistribution(
            rand, mean, stddev, NormalDistribution.DEFAULT_INVERSE_ABSOLUTE_ACCURACY);

    List<Quantile> quantiles = new ArrayList<>();
    quantiles.add(new Quantile(0.10, 0.001));
    quantiles.add(new Quantile(0.50, 0.01));
    quantiles.add(new Quantile(0.90, 0.001));
    quantiles.add(new Quantile(0.95, 0.001));
    quantiles.add(new Quantile(0.99, 0.001));

    CKMSQuantiles ckms = new CKMSQuantiles(quantiles.toArray(new Quantile[] {}));

    final int elemCount = 1000 * 1000;
    double[] shuffle = normalDistribution.sample(elemCount);

    // insert a million samples
    for (double v : shuffle) {
      ckms.insert(v);
    }

    // give the actual values for the quantiles we test
    double p10 = normalDistribution.inverseCumulativeProbability(0.1);
    double p90 = normalDistribution.inverseCumulativeProbability(0.9);
    double p95 = normalDistribution.inverseCumulativeProbability(0.95);
    double p99 = normalDistribution.inverseCumulativeProbability(0.99);

    // ε-approximate quantiles relaxes the requirement
    // to finding an item with rank between (φ−ε)n and (φ+ε)n.
    assertThat(ckms.get(0.1))
        .isCloseTo(p10, offset(errorBoundsNormalDistribution(0.1, 0.001, normalDistribution)));
    assertThat(ckms.get(0.5))
        .isCloseTo(mean, offset(errorBoundsNormalDistribution(0.5, 0.01, normalDistribution)));
    assertThat(ckms.get(0.9))
        .isCloseTo(p90, offset(errorBoundsNormalDistribution(0.9, 0.001, normalDistribution)));
    assertThat(ckms.get(0.95))
        .isCloseTo(p95, offset(errorBoundsNormalDistribution(0.95, 0.001, normalDistribution)));
    assertThat(ckms.get(0.99))
        .isCloseTo(p99, offset(errorBoundsNormalDistribution(0.99, 0.001, normalDistribution)));

    assertThat(ckms.samples).as("sample size should be below 1000").hasSizeLessThan(1000);
  }

  double errorBoundsNormalDistribution(double p, double epsilon, NormalDistribution nd) {
    // (φ+ε)n
    double upperBound = nd.inverseCumulativeProbability(p + epsilon);
    // (φ−ε)n
    double lowerBound = nd.inverseCumulativeProbability(p - epsilon);
    // subtract and divide by 2, assuming that the increase is linear in this small epsilon.
    return Math.abs(upperBound - lowerBound) / 2;
  }

  @Test
  public void testIllegalArgumentException() {
    try {
      new Quantile(-1, 0);
    } catch (IllegalArgumentException e) {
      assertThat(e.getMessage()).isEqualTo("Quantile must be between 0 and 1");
    } catch (Exception e) {
      fail("Wrong exception thrown" + e);
    }
    try {
      new Quantile(0.95, 2);
    } catch (IllegalArgumentException e) {
      assertThat(e.getMessage()).isEqualTo("Epsilon must be between 0 and 1");
    } catch (Exception e) {
      fail("Wrong exception thrown" + e);
    }
  }

  private List<Double> shuffledValues(int n, Random random) {
    List<Double> result = new ArrayList<>(n);
    for (int i = 0; i < n; i++) {
      result.add(i + 1.0);
    }
    Collections.shuffle(result, random);
    return result;
  }

  /** The following invariant must be true for each sample: g + delta <= f(r) */
  private void validateSamples(CKMSQuantiles ckms) {
    double prev = -1.0;
    int r = 0; // sum of all g's left of the current sample
    for (CKMSQuantiles.Sample sample : ckms.samples) {
      String msg =
          "invalid sample " + sample + ": count=" + ckms.n + " r=" + r + " f(r)=" + ckms.f(r);
      assertThat(sample.g + sample.delta).as(msg).isLessThanOrEqualTo(ckms.f(r));
      assertThat(prev)
          .as(
              "Samples not ordered. Keep in mind that insertBatch() takes a sorted array as"
                  + " parameter.")
          .isLessThanOrEqualTo(sample.value);
      prev = sample.value;
      r += sample.g;
    }
    assertThat(ckms.n)
        .as("the sum of all g's must be the total number of observations")
        .isEqualTo(r);
  }

  /**
   * The values that we insert in these tests are always the numbers from 1 to n, in random order.
   * So we can trivially calculate the range of acceptable results for each quantile. We check if
   * the value returned by get() is within the range of acceptable results.
   */
  private void validateResults(CKMSQuantiles ckms) {
    for (Quantile q : ckms.quantiles) {
      double actual = ckms.get(q.quantile);
      double lowerBound, upperBound;
      if (q.quantile == 0) {
        lowerBound = 1;
        upperBound = 1;
      } else if (q.quantile == 1) {
        lowerBound = ckms.n;
        upperBound = ckms.n;
      } else {
        lowerBound = Math.floor(ckms.n * (q.quantile - 2 * q.epsilon));
        upperBound = Math.ceil(ckms.n * (q.quantile + 2 * q.epsilon));
      }
      boolean ok = actual >= lowerBound && actual <= upperBound;
      if (!ok) {
        for (CKMSQuantiles.Sample sample : ckms.samples) {
          System.err.println(sample);
        }
      }
      String errorMessage =
          q
              + ": "
              + actual
              + " not in ["
              + lowerBound
              + ", "
              + upperBound
              + "], n="
              + ckms.n
              + ", "
              + q.quantile
              + "*"
              + ckms.n
              + "="
              + (q.quantile * ckms.n);
      assertThat(ok).as(errorMessage).isTrue();
    }
  }
}
