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

#include <plot/data/detail/cover_count.hpp>
#include <plot/data/detail/merge_count.hpp>
#include <plot/data/minmax_pyramid.hpp>

#include <algorithm>
#include <cmath>
#include <cstddef>

namespace plot::data {
namespace {

/// Samples covered by one level-l bucket. Clamped to kMaxPyramidLevel: the
/// shift past level 57 would overflow u64 (UB), and a bucket wider than any
/// u64 sample count is the same bucket as the widest one anyway.
constexpr u64 level_span(u32 l) {
  return u64{kBucket} << (l > kMaxPyramidLevel ? kMaxPyramidLevel : l);
}

/// Level-l buckets needed to hold n samples (ceiling division).
constexpr u64 buckets_for(u64 n, u32 l) {
  const u64 span = level_span(l);
  return n / span + (n % span == 0 ? u64{0} : u64{1});
}

/// Ceiling division: first level-l bucket whose span starts at or after v.
constexpr u64 ceil_div(u64 v, u64 d) { return v / d + (v % d == 0 ? u64{0} : u64{1}); }

/// First level-0 bucket boundary at or after v.
constexpr u64 ceil_to_bucket(u64 v) { return ceil_div(v, kBucket) * kBucket; }

/// Canonical, coarsest-first decomposition of the level-0 bucket span [b0, b1)
/// into pyramid entries, calling fn(level, absolute_bucket_index) for each.
///
/// A level-l entry is emitted only when the whole level-l bucket exists (all
/// kBucket << l samples are present, so its aggregate is final). Callers
/// guarantee b0 >= the first whole bucket of the window and b1 <= the number of
/// complete buckets, so every emission both exists and lies inside the window.
///
/// The level is chosen per *emission*, not once per pass: at each cursor the
/// largest level is taken whose bucket is aligned with the cursor, lies wholly
/// inside [b0, b1) and has already been materialised. A cursor that starts
/// misaligned (an odd first whole bucket, i.e. half of all measurement ranges)
/// therefore burns at most one level-0 bucket and then resumes the coarse
/// levels, so the decomposition is at most two entries per level and every query
/// below is O(log n) for any start. Choosing the level once per pass instead left
/// the level-0 loop to consume the whole remainder one bucket at a time, because
/// no coarser level is ever reachable from an odd cursor.
template <typename Fn>
void cover_buckets(const MinMaxPyramid& p, u64 b0, u64 b1, Fn&& fn) {
  // One past the last complete level-0 bucket: levels sit on top of those, so
  // this also bounds every coarser bucket the cursor may take.
  const u64 limit = std::min(b1, p.level_base(0) + p.level_size(0));
  const u32 top = p.level_count() - 1;         // level_count() is at least 1
  u64 b = b0;
  while (b < b1) {
    u32 li = top;
    for (;;) {
      const u64 lb = u64{1} << li;
      const u64 i = b >> li;
      if ((b % lb) == 0 && b + lb <= limit && i >= p.level_base(li) &&
          i - p.level_base(li) < p.level_size(li)) {
        break;
      }
      if (li == 0) break;
      --li;
    }
#ifdef PLOT_PYRAMID_COUNT_COVERS
    detail::cover_entry_count().fetch_add(1, std::memory_order_relaxed);
#endif
    fn(li, b >> li);
    b += u64{1} << li;
  }
}

/// Folds the decomposition of [b0, b1) into a min/max Range. Returns a
/// zero-span Range when there is nothing to fold. Buckets with no finite
/// sample (count 0) are absent, not zero-valued, and are skipped.
Range fold_range(const MinMaxPyramid& p, u64 b0, u64 b1) {
  Range r{};
  bool first = true;
  cover_buckets(p, b0, b1, [&](u32 l, u64 i) {
    if (p.bucket_count(l, i) == 0) return;
    const MinMax m = p.bucket(l, i);
    if (first) {
      r.min = f64(m.min);
      r.max = f64(m.max);
      first = false;
    } else {
      r.min = std::min(r.min, f64(m.min));
      r.max = std::max(r.max, f64(m.max));
    }
  });
  return first ? Range{0.0, 0.0} : r;
}

}  // namespace

void MinMaxPyramid::add_level() { levels_.emplace_back(); }

void MinMaxPyramid::reserve(u64 n) {
  // Materialise every level n samples can need -- level 0 plus one per doubling
  // of the bucket, capped at kMaxPyramidLevel -- so a run of appends never grows
  // the level vector. See the header for why the bucket deques themselves cannot
  // be pre-sized.
  u32 want = 1;
  for (u64 span = level_span(0); span <= n && want <= kMaxPyramidLevel; span <<= 1) {
    ++want;
  }
  if (levels_.size() < want) levels_.resize(want);
}

void MinMaxPyramid::push(u32 l, u64 k, const Entry& e) {
  Level& lv = levels_[l];
  // The first entry of a level fixes its base; every later entry is the next
  // absolute bucket, which is what keeps a level one contiguous run.
  if (lv.entries.empty()) lv.base = k;
  lv.entries.push_back(e);
#ifdef PLOT_PYRAMID_COUNT_MERGES
  detail::merge_count().fetch_add(1, std::memory_order_relaxed);
#endif
}

void MinMaxPyramid::append(f32 v) {
  // The new sample always lands in the in-progress level-0 bucket, which is the
  // only bucket any append ever touches: a bucket becomes an entry exactly once,
  // complete, and is never rewritten.
  if (count_ % kBucket == 0) pending_ = Entry{};
  // Non-finite samples occupy their absolute position -- bucket boundaries are
  // positional, so NaN or not the sample advances count_ and completes its
  // bucket -- but they are excluded from every aggregate (see the header).
  if (std::isfinite(v)) {
    const f64 vd = f64(v);
    if (pending_.count == 0) {
      pending_.mm.min = v;
      pending_.mm.max = v;
    } else {
      pending_.mm.min = std::min(pending_.mm.min, v);
      pending_.mm.max = std::max(pending_.mm.max, v);
    }
    pending_.sum += vd;
    pending_.sumsq += vd * vd;
    ++pending_.count;
  }
  ++count_;
  if (count_ % kBucket != 0) return;   // the bucket is still partial: nothing to push

  const u64 k0 = count_ / kBucket - 1;   // the level-0 bucket that just completed
  if (k0 < level_base(0)) return;        // its span starts before the window: never read
  const Entry done = pending_;
  push(0, k0, done);

  // Cascade: level l gains an entry exactly when kBucket << l divides the sample
  // count, so the walk stops at the first level that gained nothing -- amortised
  // O(1) per sample (1/kBucket + 1/(2*kBucket) + ... = 1/(kBucket/2) merges).
  // Capped at kMaxPyramidLevel: past it the span shift would overflow.
  for (u32 l = 1; l <= kMaxPyramidLevel; ++l) {
    const u64 span = level_span(l);
    if (count_ % span != 0) break;
    const u64 k = count_ / span - 1;        // the level-l bucket that just completed
    const u64 child = 2 * k;                // its left child, at level l-1
    if (child < level_base(l - 1)) break;   // a child is gone: this parent is not retained
    if (l == static_cast<u32>(levels_.size())) add_level();
    if (l >= live_levels_) live_levels_ = l + 1;
    const Entry& a = at(l - 1, child);
    const Entry& b = at(l - 1, child + 1);
    Entry m;
    // An empty child (count 0) has zero sums and a meaningless mm: it must not
    // vote 0 into the merge, so a merge with it is just the other child.
    m.mm = a.count == 0 ? b.mm : (b.count == 0 ? a.mm
                                               : MinMax{std::min(a.mm.min, b.mm.min),
                                                        std::max(a.mm.max, b.mm.max)});
    m.sum = a.sum + b.sum;
    m.sumsq = a.sumsq + b.sumsq;
    m.count = a.count + b.count;
    push(l, k, m);
  }
}

void MinMaxPyramid::drop_before(u64 absolute_sample) {
  if (absolute_sample <= base_) return;
  if (absolute_sample > count_) {
    // The producer produced more than the ring held between two acquires: every
    // appended sample is gone and the ones in between were never appended, so the
    // pyramid restarts empty at the new base rather than splicing a gap.
    base_ = absolute_sample;
    count_ = absolute_sample;
    pending_ = Entry{};
    live_levels_ = 1;
    for (u32 l = 0; l < levels_.size(); ++l) {
      levels_[l].entries.clear();
      levels_[l].base = ceil_div(absolute_sample, level_span(l));
    }
    return;
  }

  base_ = absolute_sample;
  for (u32 l = 0; l < levels_.size(); ++l) {
    Level& lv = levels_[l];
    // A bucket is dropped once its span starts before the window: the first
    // retained bucket of every level is the first one wholly inside the window,
    // which is exactly what the queries assume.
    const u64 keep = ceil_div(base_, level_span(l));
    if (keep <= lv.base) continue;
    const u64 have = static_cast<u64>(lv.entries.size());
    const u64 drop = std::min(keep - lv.base, have);
    lv.entries.erase(lv.entries.begin(),
                     lv.entries.begin() + static_cast<std::ptrdiff_t>(drop));
    // The base is where the next entry goes: past the last one if the level was
    // emptied, which is how the cursor catches up with the sliding window.
    if (drop == have) {
      lv.base = keep;
    } else {
      lv.base += drop;
    }
  }
  // The in-progress level-0 bucket may now start before the window's base: its
  // aggregate covers samples that are gone. It belongs to the bucket the window
  // starts inside, which is never read (the queries take whole buckets and the
  // Series scans the two partial edges raw), so discarding it is lossless.
  if (ceil_div(base_, kBucket) > count_ / kBucket) pending_ = Entry{};
  while (live_levels_ > 1 && level_size(live_levels_ - 1) == 0) --live_levels_;
}

void MinMaxPyramid::clear() {
  count_ = 0;
  base_ = 0;
  live_levels_ = 1;
  levels_.resize(1);
  levels_[0].entries.clear();
  levels_[0].base = 0;
  pending_ = Entry{};
}

u64 MinMaxPyramid::level_base(u32 l) const noexcept {
  return l < levels_.size() ? levels_[l].base : u64{0};
}

u64 MinMaxPyramid::level_size(u32 l) const noexcept {
  return l < levels_.size() ? static_cast<u64>(levels_[l].entries.size()) : u64{0};
}

MinMax MinMaxPyramid::bucket(u32 l, u64 i) const { return at(l, i).mm; }

f64 MinMaxPyramid::bucket_sum(u32 l, u64 i) const { return at(l, i).sum; }

f64 MinMaxPyramid::bucket_sumsq(u32 l, u64 i) const { return at(l, i).sumsq; }

u32 MinMaxPyramid::bucket_count(u32 l, u64 i) const { return at(l, i).count; }

Range MinMaxPyramid::reduce_buckets(u64 begin, u64 end) const {
  const u64 e = std::min(end, size());
  if (begin >= e) return Range{0.0, 0.0};
  // The bucket grid is absolute while these are window coordinates: window index
  // i is absolute window_base() + i, so the first whole bucket inside the range is
  // the first absolute bucket boundary at or after that.
  const u64 b0 = ceil_to_bucket(base_ + begin) / kBucket;
  const u64 b1 = (base_ + e) / kBucket;
  if (b1 <= b0) return Range{0.0, 0.0};
  return fold_range(*this, b0, b1);
}

f64 MinMaxPyramid::sum_buckets(u64 begin, u64 end) const {
  const u64 e = std::min(end, size());
  if (begin >= e) return 0.0;
  const u64 b0 = ceil_to_bucket(base_ + begin) / kBucket;
  const u64 b1 = (base_ + e) / kBucket;
  if (b1 <= b0) return 0.0;
  f64 total = 0.0;
  cover_buckets(*this, b0, b1, [&](u32 l, u64 i) { total += at(l, i).sum; });
  return total;
}

f64 MinMaxPyramid::sumsq_buckets(u64 begin, u64 end) const {
  const u64 e = std::min(end, size());
  if (begin >= e) return 0.0;
  const u64 b0 = ceil_to_bucket(base_ + begin) / kBucket;
  const u64 b1 = (base_ + e) / kBucket;
  if (b1 <= b0) return 0.0;
  f64 total = 0.0;
  cover_buckets(*this, b0, b1, [&](u32 l, u64 i) { total += at(l, i).sumsq; });
  return total;
}

u32 MinMaxPyramid::level_for(u64 n, f64 samples_per_pixel) const {
  // The ratio test also rejects NaN, and the clamp keeps a huge (or infinite)
  // samples-per-pixel from making the floor-to-u32 conversion ill-formed.
  const f64 ratio = samples_per_pixel / f64(kBucket);
  if (n == 0 || !(ratio > 1.0)) return 0;
  const f64 lg = std::floor(std::log2(ratio));
  u32 l = lg >= f64(kMaxPyramidLevel) ? kMaxPyramidLevel : static_cast<u32>(lg);
  while (l > 0) {
    const bool exists = l < level_count() && level_size(l) >= 2;
    if (exists && buckets_for(n, l) >= 2) break;
    --l;
  }
  return l;
}

u64 MinMaxPyramid::snap(f64 x, f64 x0, f64 dx) const {
  if (dx <= 0.0 || size() == 0) return 0;
  const f64 i = (x - x0) / dx;
  // Bound the round before it: llround() is undefined for values a long long
  // cannot hold, and this also maps NaN to index 0.
  if (!(i > 0.0)) return 0;
  if (i >= f64(size())) return size() - 1;
  const u64 idx = static_cast<u64>(std::llround(i));
  return idx >= size() ? size() - 1 : idx;
}

}  // namespace plot::data