// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2026 Martino Ferrari

#include "harness.hpp"
#include <plot/data/detail/cover_count.hpp>
#include <plot/data/detail/merge_count.hpp>
#include <plot/data/minmax_pyramid.hpp>
#include <algorithm>
#include <random>

using plot::data::kBucket;
using plot::data::MinMax;
using plot::data::MinMaxPyramid;
using plot::Range;

namespace {

/// Brute-force aggregate of the absolute span of level-`l` bucket `b`, from the
/// same value function the pyramid was fed. This is the oracle the absolute
/// bucket indexing is judged against: a bucket that has slid inside the window
/// must still aggregate exactly its own absolute samples.
template <typename ValueFn>
MinMax brute_bucket(u32 l, u64 b, const ValueFn& value) {
  const u64 span = u64{kBucket} << l;
  const u64 first = b * span;
  MinMax m{value(first), value(first)};
  for (u64 i = first + 1; i < first + span; ++i) {
    const f32 v = value(i);
    m.min = std::min(m.min, v);
    m.max = std::max(m.max, v);
  }
  return m;
}

}  // namespace

PLOT_TEST(pyramid_bucket_extremes_preserved) {
  MinMaxPyramid p;
  p.reserve(1000);
  for (int i = 0; i < 1000; ++i) p.append(i == 617 ? -50.0f : (i == 42 ? 900.0f : 1.0f));
  PLOT_CHECK(p.size() == 1000);
  const Range r = p.reduce_buckets(0, 960);        // exactly 15 whole buckets
  PLOT_CHECK_NEAR(r.min, -50.0, 1e-6);
  PLOT_CHECK_NEAR(r.max, 900.0, 1e-6);
}

PLOT_TEST(pyramid_reduce_buckets_matches_bruteforce_on_aligned_ranges) {
  std::mt19937 rng(1234);
  std::uniform_real_distribution<float> d(-100.0f, 100.0f);
  std::vector<float> v(5000);
  MinMaxPyramid p;
  for (auto& x : v) { x = d(rng); p.append(x); }
  for (u64 a = 0; a < 5000; a += kBucket * 3) {
    for (u64 b = a + kBucket; b <= 5000; b += kBucket * 7) {
      const u64 ahi = (b / kBucket) * kBucket;     // whole buckets only
      if (ahi <= a) continue;
      float lo = v[a], hi = v[a];
      for (u64 i = a; i < ahi; ++i) { lo = std::min(lo, v[i]); hi = std::max(hi, v[i]); }
      const Range r = p.reduce_buckets(a, ahi);
      PLOT_CHECK_NEAR(r.min, lo, 1e-4);
      PLOT_CHECK_NEAR(r.max, hi, 1e-4);
    }
  }
}

PLOT_TEST(pyramid_reduce_buckets_excludes_partial_edges) {
  // Bucket 1 covers samples [64,128). reduce_buckets answers in whole buckets
  // only, so the spike must show up in its own bucket and nowhere else.
  MinMaxPyramid p;
  for (int i = 0; i < 128; ++i) p.append(i == 70 ? 999.0f : 1.0f);
  PLOT_CHECK_NEAR(p.reduce_buckets(64, 128).max, 999.0, 1e-6);
  PLOT_CHECK_NEAR(p.reduce_buckets(0, 64).max, 1.0, 1e-6);
  PLOT_CHECK(p.reduce_buckets(64, 64).span() == 0.0);   // empty -> degenerate
}

PLOT_TEST(pyramid_aligned_queries_never_include_a_partial_bucket) {
  // Corpus of 100 samples: bucket 0 is whole, bucket 1 (samples 64..127) is
  // partial. A whole-bucket query must not reach into the partial one, in
  // either direction.
  MinMaxPyramid p;
  for (int i = 0; i < 100; ++i) p.append(i == 99 ? 777.0f : 1.0f);
  PLOT_CHECK(p.reduce_buckets(64, 100).span() == 0.0);   // no whole bucket inside
  PLOT_CHECK_NEAR(p.reduce_buckets(0, 100).max, 1.0, 1e-6);   // bucket 0 only
  PLOT_CHECK(p.sum_buckets(64, 100) == 0.0);
  PLOT_CHECK(p.sumsq_buckets(64, 100) == 0.0);
}

PLOT_TEST(pyramid_level_selection) {
  MinMaxPyramid p;
  p.reserve(1u << 20);
  for (u64 i = 0; i < (1u << 20); ++i) p.append(0.0f);
  // <= kBucket samples per pixel must use level 0 (raw path handles it anyway).
  PLOT_CHECK(p.level_for(1u << 20, 1.0) == 0);
  // 512 samples per pixel with kBucket == 64 → bucket size 512 → level 3.
  PLOT_CHECK(p.level_for(1u << 20, 512.0) == 3);
  // Level selection must never exceed the deepest level that has data.
  PLOT_CHECK(p.level_for(1000, 1e9) <= p.level_count() - 1);
}

PLOT_TEST(pyramid_snap_returns_nearest_sample) {
  MinMaxPyramid p;
  for (int i = 0; i < 100; ++i) p.append(float(i));
  PLOT_CHECK(p.snap(10.4, 0.0, 1.0) == 10);
  PLOT_CHECK(p.snap(10.6, 0.0, 1.0) == 11);
  PLOT_CHECK(p.snap(-5.0, 0.0, 1.0) == 0);      // clamped
  PLOT_CHECK(p.snap(1e9, 0.0, 1.0) == 99);      // clamped
}

PLOT_TEST(pyramid_sums_match_bruteforce_on_aligned_ranges) {
  std::mt19937 rng(99);
  std::uniform_real_distribution<float> d(-3.0f, 7.0f);
  std::vector<float> v(3000);
  MinMaxPyramid p;
  for (auto& x : v) { x = d(rng); p.append(x); }
  for (u64 a = 0; a < 3000; a += kBucket * 5) {
    for (u64 b = a + kBucket; b <= 3000; b += kBucket * 9) {
      double s = 0.0, sq = 0.0;
      for (u64 i = a; i < b; ++i) { s += v[i]; sq += double(v[i]) * double(v[i]); }
      PLOT_CHECK_NEAR(p.sum_buckets(a, b), s, 1e-3);
      PLOT_CHECK_NEAR(p.sumsq_buckets(a, b), sq, 1e-2);
    }
  }
}

PLOT_TEST(pyramid_clear_and_rebuild) {
  MinMaxPyramid p;
  for (int i = 0; i < 500; ++i) p.append(1.0f);
  p.clear();
  PLOT_CHECK(p.size() == 0);
  PLOT_CHECK(p.level_count() <= 1);
  p.append(7.0f);
  PLOT_CHECK(p.size() == 1);
  // A pyramid can only answer exactly for whole buckets, so complete the first
  // bucket and ask the bucket-grained query for its extrema.
  for (u32 i = 1; i < kBucket; ++i) p.append(1.0f);
  PLOT_CHECK(p.size() == kBucket);
  PLOT_CHECK(p.level_base(0) == 0);
  PLOT_CHECK(p.level_size(0) == 1);
  PLOT_CHECK_NEAR(p.reduce_buckets(0, kBucket).max, 7.0, 1e-6);
  PLOT_CHECK_NEAR(p.reduce_buckets(0, kBucket).min, 1.0, 1e-6);
}

PLOT_TEST(pyramid_cover_cost_is_logarithmic_for_unaligned_starts) {
  // Answers cannot show this: a decomposition that walks the range one level-0
  // bucket at a time returns exactly the same min/max as the O(log n) one. So
  // count the folded aggregates. For this query the bad pass emitted 15624
  // entries (the level-0 loop consumed the whole remainder); picking a level per
  // emission leaves 19, at most two per level, and the aligned start stays at 7.
  MinMaxPyramid p;
  p.reserve(1'000'000);
  for (u64 i = 0; i < 1'000'000; ++i) p.append(0.0f);
  PLOT_CHECK(p.level_size(0) == 1'000'000 / kBucket);   // 15625, no partial bucket

  auto folds_for = [&](u64 begin, u64 end) {
    plot::data::detail::cover_entry_count().store(0, std::memory_order_relaxed);
    (void)p.reduce_buckets(begin, end);
    return plot::data::detail::cover_entry_count().load(std::memory_order_relaxed);
  };
  const u64 aligned = folds_for(0, 1'000'000);     // first whole bucket 0
  const u64 unaligned = folds_for(1, 1'000'000);   // first whole bucket 1: odd
  PLOT_CHECK(aligned <= 64);
  PLOT_CHECK(unaligned < 100);
  PLOT_CHECK(unaligned <= u64{4} * aligned);       // misalignment costs entries, not a scan
}

PLOT_TEST(pyramid_drop_before_slides_the_window) {
  MinMaxPyramid p;
  for (u64 i = 0; i < 4096; ++i) p.append(f32(i % 13));
  PLOT_CHECK(p.window_base() == 0);
  PLOT_CHECK(p.window_end() == 4096);
  PLOT_CHECK(p.size() == 4096);

  p.drop_before(1000);
  PLOT_CHECK(p.window_base() == 1000);
  PLOT_CHECK(p.size() == 3096);
  // The window's first *whole* bucket starts at 1024: the bucket the drop cut
  // into is not representable as a whole bucket and is never queried.
  PLOT_CHECK(p.level_base(0) == 1024 / kBucket);
  PLOT_CHECK(p.level_base(0) + p.level_size(0) == 4096 / kBucket);

  f32 lo = 1.0f, hi = -1.0f;
  for (u64 i = 1024; i < 4096; ++i) {
    lo = std::min(lo, f32(i % 13));
    hi = std::max(hi, f32(i % 13));
  }
  const Range r = p.reduce_buckets(0, p.size());
  PLOT_CHECK_NEAR(r.min, lo, 1e-6);
  PLOT_CHECK_NEAR(r.max, hi, 1e-6);
  // Window [1000, 1024) holds no whole bucket at all: an empty answer, not one
  // that quietly reaches into the dropped samples.
  PLOT_CHECK(p.reduce_buckets(0, 24).span() == 0.0);
  PLOT_CHECK(p.sum_buckets(0, 24) == 0.0);
  PLOT_CHECK(p.sumsq_buckets(0, 24) == 0.0);

  p.drop_before(1000);                    // at the base: no-op
  PLOT_CHECK(p.window_base() == 1000);
  p.drop_before(4096);                    // at the end: the window is empty
  PLOT_CHECK(p.size() == 0);
  p.append(5.0f);                         // and appending starts a fresh window
  PLOT_CHECK(p.size() == 1);
  PLOT_CHECK(p.window_base() == 4096);
}

PLOT_TEST(pyramid_bucket_aggregates_stay_absolute_across_slides) {
  // The invariant the sliding window rests on: a bucket that has slid inside the
  // window is never recomputed, so it must already hold the aggregate of its own
  // absolute span -- at every level, after any number of drops.
  MinMaxPyramid p;
  const auto value = [](u64 i) { return f32((i * 37u) % 101u) - 50.0f; };
  constexpr u64 kTotal = 20'000;
  constexpr u64 kWindow = 4'000;
  for (u64 i = 0; i < kTotal; ++i) {
    p.append(value(i));
    if ((i + 1) % 997 == 0 && i + 1 > kWindow) p.drop_before(i + 1 - kWindow);
  }
  // The last drop is the last multiple of 997 at or below kTotal.
  PLOT_CHECK(p.window_base() == (kTotal / 997) * 997 - kWindow);
  for (u32 l = 0; l < p.level_count(); ++l) {
    PLOT_CHECK(p.level_size(l) > 0);
    const u64 span = u64{kBucket} << l;
    for (u64 b = p.level_base(l); b < p.level_base(l) + p.level_size(l); ++b) {
      const MinMax want = brute_bucket(l, b, value);
      const MinMax got = p.bucket(l, b);
      PLOT_CHECK_NEAR(got.min, want.min, 1e-6);
      PLOT_CHECK_NEAR(got.max, want.max, 1e-6);
      PLOT_CHECK(b * span >= p.window_base());          // wholly inside the window
      PLOT_CHECK((b + 1) * span <= p.window_end());     // and complete
    }
  }
}

PLOT_TEST(pyramid_sliding_window_bucket_queries_match_bruteforce) {
  MinMaxPyramid p;
  const auto value = [](u64 i) { return f32(i % 89) * 0.25f - 5.0f; };
  constexpr u64 kTotal = 12'000;
  constexpr u64 kWindow = 3'000;
  for (u64 i = 0; i < kTotal; ++i) {
    p.append(value(i));
    if ((i + 1) % 613 != 0) continue;
    const u64 abs0 = i + 1 > kWindow ? i + 1 - kWindow : 0;
    p.drop_before(abs0);
    const u64 n = p.size();
    // Window coordinates: window index j is absolute abs0 + j.
    const u64 first = abs0 / kBucket + (abs0 % kBucket == 0 ? u64{0} : u64{1});
    const u64 last = (abs0 + n) / kBucket;
    if (last <= first) continue;
    f32 lo = 1e30f, hi = -1e30f;
    f64 sum = 0.0, sq = 0.0;
    for (u64 a = first * kBucket; a < last * kBucket; ++a) {
      const f32 v = value(a);
      lo = std::min(lo, v);
      hi = std::max(hi, v);
      sum += f64(v);
      sq += f64(v) * f64(v);
    }
    const u64 wb = first * kBucket - abs0;
    const u64 we = last * kBucket - abs0;
    const Range r = p.reduce_buckets(wb, we);
    PLOT_CHECK_NEAR(r.min, lo, 1e-6);
    PLOT_CHECK_NEAR(r.max, hi, 1e-6);
    PLOT_CHECK_NEAR(p.sum_buckets(wb, we), sum, 1e-3);
    PLOT_CHECK_NEAR(p.sumsq_buckets(wb, we), sq, 1e-1);
  }
}

PLOT_TEST(pyramid_append_cost_is_amortised_constant) {
  // A sample must not walk the live levels: level l gains an entry once per
  // kBucket << l samples, so N appends perform ~N/32 merges however many levels
  // the pyramid holds. Walking every level per append costs ~N * levels.
  MinMaxPyramid p;
  p.reserve(1u << 20);
  constexpr u64 kAppends = 1u << 20;
  for (u64 i = 0; i < kAppends; ++i) p.append(0.5f);
  PLOT_CHECK(p.level_count() >= 10);   // 2^20 samples: 16384 level-0 buckets

  plot::data::detail::merge_count().store(0, std::memory_order_relaxed);
  for (u64 i = 0; i < kAppends; ++i) p.append(0.5f);
  const u64 merges = plot::data::detail::merge_count().load(std::memory_order_relaxed);
  // ~N/64 + N/128 + ... = ~N/32; the bound leaves 2x room, and is four orders of
  // magnitude below the N * live_levels an unbounded walk would perform.
  PLOT_CHECK(merges <= kAppends / 16);

  // Sliding must not add per-sample work either: the appends below run against a
  // window whose base moves with every sample.
  plot::data::detail::merge_count().store(0, std::memory_order_relaxed);
  for (u64 i = 0; i < kAppends; ++i) {
    p.drop_before(i + 1);
    p.append(0.5f);
  }
  const u64 slid = plot::data::detail::merge_count().load(std::memory_order_relaxed);
  PLOT_CHECK(slid <= kAppends / 16);
}