#include "harness.hpp"
#include <plot/data/array_series.hpp>
#include <plot/data/detail/merge_count.hpp>
#include <plot/data/stream_series.hpp>
#include <algorithm>
#include <array>
#include <atomic>
#include <chrono>
#include <cmath>
#include <random>
#include <thread>
#include <vector>

using plot::data::kBucket;

PLOT_TEST(array_series_bounds_and_gather) {
  std::vector<f32> y{1, 5, -2, 3};
  plot::data::ArraySeries s(std::move(y), /*x0=*/10.0, /*dx=*/0.5);
  PLOT_CHECK(s.size() == 4);
  PLOT_CHECK_NEAR(s.x_bounds().min, 10.0, 1e-9);
  PLOT_CHECK_NEAR(s.x_bounds().max, 11.5, 1e-9);
  PLOT_CHECK_NEAR(s.y_bounds().min, -2.0, 1e-6);
  std::vector<f32> out;
  s.gather(0, 0, 4, out);
  PLOT_CHECK(out.size() == 8);             // interleaved x,y
  PLOT_CHECK_NEAR(out[0], 10.0, 1e-6);
  PLOT_CHECK_NEAR(out[1], 1.0, 1e-6);
  PLOT_CHECK_NEAR(out[6], 11.5, 1e-6);
}

PLOT_TEST(stream_series_acquire_publishes_snapshot) {
  plot::data::StreamSeries s(1 << 12, /*x0=*/0.0, /*dx=*/1e-6);
  for (int i = 0; i < 100; ++i) s.push(f32(i));
  PLOT_CHECK(s.produced() == 100);
  s.acquire();
  PLOT_CHECK(s.size() == 100);
  // Pyramid must track the published samples.
  PLOT_CHECK_NEAR(s.reduce(0, 100).max, 99.0, 1e-4);
  for (int i = 100; i < 250; ++i) s.push(f32(i));
  s.acquire();
  PLOT_CHECK(s.size() == 250);
}

PLOT_TEST(stream_series_drops_oldest_when_full) {
  plot::data::StreamSeries s(256, 0.0, 1.0);
  for (int i = 0; i < 1000; ++i) s.push(f32(i));
  s.acquire();
  PLOT_CHECK(s.size() <= 256);
  PLOT_CHECK_NEAR(s.y_bounds().max, 999.0, 1e-4);
}

PLOT_TEST(array_series_exact_range_queries_match_bruteforce) {
  // The strong test for the exactness split: arbitrary, deliberately
  // unaligned ranges must match a brute-force scan of the raw samples.
  std::mt19937 rng(7);
  std::uniform_real_distribution<float> d(-5.0f, 5.0f);
  std::vector<f32> y(4000);
  for (auto& v : y) v = d(rng);
  plot::data::ArraySeries s(y, 0.0, 1.0);
  for (u64 a = 0; a < 4000; a += 173) {
    for (u64 b = a + 1; b <= 4000; b += 331) {
      float lo = y[a], hi = y[a];
      double sum = 0.0, sq = 0.0;
      for (u64 i = a; i < b; ++i) {
        lo = std::min(lo, y[i]); hi = std::max(hi, y[i]);
        sum += y[i]; sq += double(y[i]) * double(y[i]);
      }
      const plot::Range r = s.reduce(a, b);
      PLOT_CHECK_NEAR(r.min, lo, 1e-4);
      PLOT_CHECK_NEAR(r.max, hi, 1e-4);
      PLOT_CHECK_NEAR(s.sum(a, b), sum, 1e-2);
      PLOT_CHECK_NEAR(s.sumsq(a, b), sq, 1e-1);
    }
  }
}

PLOT_TEST(stream_series_exact_queries_survive_ring_wrap) {
  plot::data::StreamSeries s(1u << 10, 0.0, 1.0);
  constexpr u64 kPushes = 3000;
  constexpr u64 kFirstRetained = kPushes - (1u << 10);  // 1976
  for (u64 i = 0; i < kPushes; ++i) s.push(f32(i % 97));   // wraps the ring twice
  s.acquire();
  const u64 n = s.size();
  PLOT_CHECK(n == 1u << 10);
  PLOT_CHECK_NEAR(s.x0(), 1976.0, 1e-9);
  PLOT_CHECK(s.x_bounds().min == s.x0());
  PLOT_CHECK(s.pyramid().size() == n);

  const plot::Range r = s.reduce(0, n);                 // retained window, wrapped
  PLOT_CHECK_NEAR(r.max, 96.0, 1e-4);
  PLOT_CHECK_NEAR(r.min, 0.0, 1e-4);

  double sum = 0.0, sumsq = 0.0;
  for (u64 i = kFirstRetained; i < kPushes; ++i) {
    const double v = static_cast<double>(i % 97);
    sum += v;
    sumsq += v * v;
  }
  PLOT_CHECK_NEAR(s.sum(0, n), sum, 1e-2);
  PLOT_CHECK_NEAR(s.sumsq(0, n), sumsq, 1e-1);

  std::vector<f32> out;
  s.gather(0, 0, n, out);                               // level 0: raw samples
  PLOT_CHECK(out.size() == 2 * n);
  PLOT_CHECK_NEAR(out[0], s.x0(), 1e-6);
  PLOT_CHECK_NEAR(out[1], f32(kFirstRetained % 97), 1e-6);
  PLOT_CHECK_NEAR(out[2 * n - 2], s.x0() + f64(n - 1), 1e-6);
}

PLOT_TEST(stream_series_push_and_read_are_race_free) {
  // The ring overwrites, so the producer reuses the slot of the snapshot's oldest
  // retained sample while the consumer may still be reading it. Relaxed atomic
  // slots make that defined; nothing else in the suite has two threads, so this is
  // the case ThreadSanitizer actually flags when the slots are plain.
  //
  // The two flags make the overlap structural, not scheduling luck: the reader
  // waits until the producer is already pushing, acknowledges, and only then does
  // the producer push the rest -- so the reader's window spans pushes that are in
  // flight by construction. The reader then loops until the producer reports done,
  // which also makes "at least one read" true by construction instead of by
  // winning a race with a fast producer.
  plot::data::StreamSeries s(1u << 8, 0.0, 1.0);   // one read spans reused slots
  constexpr u64 kStride = 512;                     // pushes before the reader starts
  constexpr u64 kPushes = 4'096;                   // total: a few thousand only
  std::atomic<bool> started{false};
  std::atomic<bool> reading{false};
  std::atomic<bool> done{false};

  std::thread producer([&] {
    for (u64 i = 0; i < kStride; ++i) s.push(f32(i % 971));
    started.store(true, std::memory_order_release);
    while (!reading.load(std::memory_order_acquire)) {
    }
    for (u64 i = kStride; i < kPushes; ++i) s.push(f32(i % 971));
    done.store(true, std::memory_order_release);
  });

  while (!started.load(std::memory_order_acquire)) {
  }
  reading.store(true, std::memory_order_release);

  u64 reads = 0;
  bool sizes_ok = true;
  bool finished = false;
  do {
    finished = done.load(std::memory_order_acquire);
    s.acquire();
    const u64 n = s.size();
    if (n > 0) {
      (void)s.reduce(0, n);
      (void)s.sum(0, n);
      (void)s.sumsq(0, n);
      std::vector<f32> out;
      s.gather(0, 0, n, out);                           // reads the oldest edge
      if (out.size() != 2 * n) sizes_ok = false;
    }
    ++reads;
  } while (!finished);

  producer.join();
  // Checked only after the join: a throw while the producer is joinable would
  // std::terminate instead of reporting the failure.
  PLOT_CHECK(s.produced() == kPushes);
  PLOT_CHECK(reads >= 1);
  PLOT_CHECK(sizes_ok);
}

namespace {

/// Buckets the pyramid holds across every level: its footprint, not its window.
u64 pyramid_entries(const plot::data::MinMaxPyramid& p) {
  u64 total = 0;
  for (u32 l = 0; l < p.level_count(); ++l) total += p.level_size(l);
  return total;
}

}  // namespace

PLOT_TEST(stream_series_sliding_window_queries_match_bruteforce) {
  // After the base has moved many times, an exact query must still match a scan
  // of exactly the retained samples. The pyramid's buckets live in absolute
  // sample space while the window's base is arbitrary, and that meeting point is
  // where an off-by-one hides -- so test the sliding case, not just the static
  // ArraySeries one.
  constexpr u64 kCapacity = 1u << 14;
  constexpr u64 kPushes = 60'000;
  const auto value = [](u64 i) { return f32((i * 29u) % 173u) - 86.0f; };
  plot::data::StreamSeries s(kCapacity, 0.0, 1e-6);
  for (u64 i = 0; i < kPushes; i += 337) {
    const u64 n = std::min<u64>(337, kPushes - i);
    for (u64 k = 0; k < n; ++k) s.push(value(i + k));
    s.acquire();   // the window's base moves on every one of these
  }

  const u64 n = s.size();
  const u64 base = s.produced() - n;
  PLOT_CHECK(n == kCapacity);
  PLOT_CHECK(base > 2 * kCapacity);             // the window really did slide
  PLOT_CHECK(s.pyramid().size() == n);          // pyramid() indexes the window
  PLOT_CHECK(s.pyramid().window_base() == base);
  PLOT_CHECK(s.pyramid().window_end() == s.produced());
  PLOT_CHECK_NEAR(s.x0(), f64(base) * 1e-6, 1e-12);
  PLOT_CHECK(s.x_bounds().min == s.x0());
  PLOT_CHECK_NEAR(s.x_bounds().max, f64(base + n - 1) * 1e-6, 1e-9);

  // Brute force in window coordinates over the retained samples only.
  const auto brute = [&](u64 a, u64 b) {
    f64 lo = f64(value(base + a)), hi = lo;
    f64 sum = 0.0, sq = 0.0;
    for (u64 i = a; i < b; ++i) {
      const f64 v = f64(value(base + i));
      lo = std::min(lo, v);
      hi = std::max(hi, v);
      sum += v;
      sq += v * v;
    }
    return std::array<f64, 4>{lo, hi, sum, sq};
  };

  const u64 starts[] = {0, 1, 63, 64, 65, 1000, 7919, n - 200};
  for (const u64 a : starts) {
    for (u64 b = a + 1; b <= n; b += 1301) {
      const std::array<f64, 4> want = brute(a, b);
      const plot::Range r = s.reduce(a, b);
      PLOT_CHECK_NEAR(r.min, want[0], 1e-4);
      PLOT_CHECK_NEAR(r.max, want[1], 1e-4);
      PLOT_CHECK_NEAR(s.sum(a, b), want[2], 1e-2);
      PLOT_CHECK_NEAR(s.sumsq(a, b), want[3], 1e-1);
    }
  }
  const std::array<f64, 4> whole = brute(0, n);
  PLOT_CHECK_NEAR(s.y_bounds().min, whole[0], 1e-6);
  PLOT_CHECK_NEAR(s.y_bounds().max, whole[1], 1e-6);

  // The level accessors are in absolute bucket coordinates: a stored bucket must
  // aggregate its own absolute span, the sliding window included.
  const u64 span = u64{kBucket} << 1;
  const u64 lbase = s.pyramid().level_base(1);
  const u64 lsize = s.pyramid().level_size(1);
  PLOT_CHECK(lsize >= 2);
  for (u64 b = lbase; b < lbase + std::min<u64>(lsize, 4); ++b) {
    const std::array<f64, 4> want = brute(b * span - base, (b + 1) * span - base);
    const plot::data::MinMax m = s.pyramid().bucket(1, b);
    PLOT_CHECK_NEAR(m.min, want[0], 1e-6);
    PLOT_CHECK_NEAR(m.max, want[1], 1e-6);
  }

  // level-1 gather: one envelope pair per retained bucket the range overlaps, x
  // at the bucket's own first and last sample. This is the mapping the renderer
  // draws from.
  std::vector<f32> out;
  s.gather(1, 0, n, out);
  const u64 first = std::max(base / span, lbase);
  const u64 last = std::min((base + n - 1) / span + 1, lbase + lsize);
  PLOT_CHECK(out.size() == 4 * (last - first));
  for (u64 bi = first; bi < last; ++bi) {
    const std::array<f64, 4> want = brute(bi * span - base, (bi + 1) * span - base);
    const plot::usize k = static_cast<plot::usize>(bi - first);
    const u64 first_sample = bi * span - base;
    const u64 last_sample = std::min((bi + 1) * span, base + n) - 1 - base;
    // x is in data coordinates: x0() is the retained window's first sample, so
    // window index j plots at (base + j) * dx.
    PLOT_CHECK_NEAR(out[4 * k], f32(f64(base + first_sample) * 1e-6), 1e-9);
    PLOT_CHECK_NEAR(out[4 * k + 1], want[0], 1e-6);
    PLOT_CHECK_NEAR(out[4 * k + 2], f32(f64(base + last_sample) * 1e-6), 1e-9);
    PLOT_CHECK_NEAR(out[4 * k + 3], want[1], 1e-6);
  }
}

PLOT_TEST(stream_series_steady_state_acquire_is_proportional_to_produced) {
  // The defect this pins: a saturated ring whose base moved every frame used to
  // rebuild the whole window, costing O(capacity * live_levels) merges per
  // acquire. It must cost O(samples produced since the last acquire) instead, so
  // the same production through a 16x larger window must cost the same.
  const auto merges_for = [](u64 capacity) {
    plot::data::StreamSeries s(capacity, 0.0, 1.0);
    constexpr u64 kSaturate = 1u << 20;
    for (u64 i = 0; i < kSaturate; ++i) s.push(f32(i % 977));
    s.acquire();
    PLOT_CHECK(s.size() == capacity);

    constexpr u64 kFrames = 64;
    constexpr u64 kFrame = 1u << 13;   // 8192 samples: 1 MS/s at ~122 fps
    plot::data::detail::merge_count().store(0, std::memory_order_relaxed);
    for (u64 f = 0; f < kFrames; ++f) {
      for (u64 k = 0; k < kFrame; ++k) s.push(f32((kSaturate + f * kFrame + k) % 977));
      s.acquire();
    }
    return plot::data::detail::merge_count().load(std::memory_order_relaxed);
  };

  constexpr u64 kProduced = 64 * (1u << 13);
  const u64 small = merges_for(1u << 16);   // 64k-sample window
  const u64 big = merges_for(1u << 20);     // 1M-sample window: 16x to rebuild
  // ~kProduced/32 in practice; rebuilding the 1M window alone would merge
  // kCapacity/64 * live_levels ~ 245'000 buckets per acquire.
  PLOT_CHECK(small <= kProduced / 16);
  PLOT_CHECK(big <= kProduced / 16);
  PLOT_CHECK(big <= 2 * small + kBucket);
}

PLOT_TEST(stream_series_steady_state_acquire_is_fast) {
  // Wall-clock backstop for the merge count above: rebuilding the 1M-sample
  // window measured ~44 ms, so 128 such frames would need ~5.6 s. The sliding
  // path does the same work in a few milliseconds; the bound is two orders of
  // magnitude clear of both, so it is not a thin timing margin.
  plot::data::StreamSeries s(1u << 20, 0.0, 1.0);
  constexpr u64 kSaturate = 1u << 20;
  for (u64 i = 0; i < kSaturate; ++i) s.push(f32(i % 977));
  s.acquire();

  constexpr u64 kFrames = 128;
  constexpr u64 kFrame = 1u << 13;
  const auto t0 = std::chrono::steady_clock::now();
  for (u64 f = 0; f < kFrames; ++f) {
    for (u64 k = 0; k < kFrame; ++k) s.push(f32((kSaturate + f * kFrame + k) % 977));
    s.acquire();
  }
  const f64 ms =
      std::chrono::duration<f64>(std::chrono::steady_clock::now() - t0).count() * 1.0e3;
  PLOT_CHECK(s.size() == 1u << 20);
  PLOT_CHECK(ms < 250.0);
}

PLOT_TEST(stream_series_pyramid_memory_is_bounded_by_the_window) {
  // 50M samples through a 1M-sample ring: the pyramid's footprint must track the
  // retained window (n/64 level-0 buckets plus the levels over them), not the
  // samples ever pushed.
  constexpr u64 kCapacity = 1u << 20;
  constexpr u64 kTotal = 50u * (1u << 20);
  constexpr u64 kChunk = 1u << 19;
  plot::data::StreamSeries s(kCapacity, 0.0, 1.0);
  u64 entries_after_10m = 0;
  for (u64 p = 0; p < kTotal; p += kChunk) {
    for (u64 i = 0; i < kChunk; ++i) s.push(f32((p + i) % 251));
    s.acquire();
    if (p + kChunk == 10u * (1u << 20)) entries_after_10m = pyramid_entries(s.pyramid());
  }
  const u64 entries = pyramid_entries(s.pyramid());
  const u64 bound = 2 * (kCapacity / kBucket) + kBucket;
  PLOT_CHECK(s.produced() == kTotal);
  PLOT_CHECK(s.size() == kCapacity);
  PLOT_CHECK(entries <= bound);
  // 5x the samples pushed, same footprint: levels deeper than one chunk can
  // differ by a bucket each, hence the slack.
  PLOT_CHECK(entries <= entries_after_10m + kBucket);
}