// Benchmark: build throughput, recall@10 and query throughput of the public
// API against a brute-force ground truth on uniform random data — for every
// distance metric, plus a larger scale run.
//
// Run with:  moon run --release examples/benchmark

///|
priv struct Rng {
  mut state : UInt64
}

///|
fn Rng::next_f64(self : Rng) -> Double {
  let state = self.state + 0x9E3779B97F4A7C15UL
  self.state = state
  let mut z = state
  z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9UL
  z = (z ^ (z >> 27)) * 0x94D049BB133111EBUL
  ((z ^ (z >> 31)) >> 11).to_double() * (1.0 / 9007199254740992.0)
}

///|
fn gen_dataset(seed : UInt64, n : Int32, dim : Int32) -> FixedArray[Float] {
  let rng = { state: seed, }
  let data : FixedArray[Float] = FixedArray::make(n * dim, 0.0)
  let mut i = 0
  while i < n * dim {
    data[i] = Float::from_double(rng.next_f64())
    i += 1
  }
  data
}

///|
fn row_of(
  data : FixedArray[Float],
  row : Int32,
  dim : Int32,
) -> FixedArray[Float] {
  let out : FixedArray[Float] = FixedArray::make(dim, 0.0)
  FixedArray::unsafe_blit(out, 0, data, row * dim, dim)
  out
}

///|
/// L2-normalizes every row of a flat dataset into a new array.
fn normalize_flat(
  data : FixedArray[Float],
  n : Int32,
  dim : Int32,
) -> FixedArray[Float] {
  let out : FixedArray[Float] = FixedArray::make(n * dim, 0.0)
  let mut i = 0
  while i < n {
    let off = i * dim
    let mut norm_sq = 0.0
    let mut j = 0
    while j < dim {
      let v = data.unsafe_get(off + j).to_double()
      norm_sq += v * v
      j += 1
    }
    j = 0
    if norm_sq > 0.0 {
      let inv = 1.0 / norm_sq.sqrt()
      while j < dim {
        out[off + j] = Float::from_double(
          data.unsafe_get(off + j).to_double() * inv,
        )
        j += 1
      }
    } else {
      while j < dim {
        out[off + j] = data.unsafe_get(off + j)
        j += 1
      }
    }
    i += 1
  }
  out
}

///|
/// Distance between query `q` and flat-dataset row `row`, allocation-free.
/// For `InnerProduct` the value is negated (smaller = more similar), which
/// also covers Cosine when both sides are pre-normalized.
fn dist_row(
  metric : @moonhnsw.Distance,
  data : FixedArray[Float],
  row : Int32,
  dim : Int32,
  q : FixedArray[Float],
) -> Double {
  let off = row * dim
  match metric {
    @moonhnsw.Distance::L2 => {
      let mut s : Float = 0.0
      let mut j = 0
      while j < dim {
        let d = data.unsafe_get(off + j) - q.unsafe_get(j)
        s += d * d
        j += 1
      }
      s.to_double()
    }
    _ => {
      let mut s : Float = 0.0
      let mut j = 0
      while j < dim {
        s += data.unsafe_get(off + j) * q.unsafe_get(j)
        j += 1
      }
      -s.to_double()
    }
  }
}

///|
/// Exact top-k ids via a single pass over all rows, maintaining a running
/// top-k — the fair brute-force baseline (n distances per query).
fn brute_topk(
  metric : @moonhnsw.Distance,
  truth_data : FixedArray[Float],
  n : Int32,
  dim : Int32,
  q : FixedArray[Float],
  k : Int32,
) -> Array[Int32] {
  // Running top-k as parallel id/distance arrays, worst at the end.
  let ids : FixedArray[Int32] = FixedArray::make(k, -1)
  let ds : FixedArray[Double] = FixedArray::make(k, 1.0e308)
  let mut filled = 0
  for i in 0.. 0 && ds.unsafe_get(j - 1) > d {
        ids[j] = ids.unsafe_get(j - 1)
        ds[j] = ds.unsafe_get(j - 1)
        j -= 1
      }
      ids[j] = i
      ds[j] = d
      filled += 1
    } else if d < ds.unsafe_get(k - 1) {
      let mut j = k - 1
      while j > 0 && ds.unsafe_get(j - 1) > d {
        ids[j] = ids.unsafe_get(j - 1)
        ds[j] = ds.unsafe_get(j - 1)
        j -= 1
      }
      ids[j] = i
      ds[j] = d
    }
  }
  let out : Array[Int32] = []
  for j in 0.. String {
  ((x * 100.0 + 0.5).to_int().to_double() / 100.0).to_string()
}

///|
fn bench_metric(
  name : String,
  metric : @moonhnsw.Distance,
  n : Int32,
  dim : Int32,
  n_queries : Int32,
  k : Int32,
  m : Int32,
  ef_construction : Int32,
  seed : UInt64,
  storage~ : @moonhnsw.Storage = @moonhnsw.Storage::Float32,
) -> Unit raise @moonhnsw.HnswError {
  println(
    "--- \{name}  n=\{n} dim=\{dim} queries=\{n_queries} k=\{k} m=\{m} ef_construction=\{ef_construction} ---",
  )
  let raw = gen_dataset(seed, n, dim)
  let queries = gen_dataset(seed + 1UL, n_queries, dim)
  // Ground truth for Cosine = negated inner product over pre-normalized
  // vectors; the index normalizes internally, the truth side does it here.
  let for_cosine = match metric {
    @moonhnsw.Distance::Cosine => true
    _ => false
  }
  let truth_metric = if for_cosine {
    @moonhnsw.Distance::InnerProduct
  } else {
    metric
  }
  let truth_data = if for_cosine { normalize_flat(raw, n, dim) } else { raw }
  // Build.
  let t0 = @bench.monotonic_clock_start()
  let idx = @moonhnsw.Hnsw::new(
    dim,
    metric,
    params=@moonhnsw.HnswParams::new(m~, ef_construction~, storage~),
  )
  for i in 0.. Unit raise @moonhnsw.HnswError {
  let dim = 32
  let n = 2000
  let rounds = 20
  let raw = gen_dataset(20261029UL, n, dim)
  let vectors : Array[FixedArray[Float]] = []
  for i in 0.. Unit raise @moonhnsw.HnswError {
  let dim = 32
  let k = 10
  let m = 16
  let ef_construction = 200
  bench_metric(
    "L2",
    @moonhnsw.Distance::L2,
    2000,
    dim,
    100,
    k,
    m,
    ef_construction,
    20261024UL,
  )
  bench_metric(
    "InnerProduct",
    @moonhnsw.Distance::InnerProduct,
    2000,
    dim,
    100,
    k,
    m,
    ef_construction,
    20261025UL,
  )
  bench_metric(
    "Cosine",
    @moonhnsw.Distance::Cosine,
    2000,
    dim,
    100,
    k,
    m,
    ef_construction,
    20261026UL,
  )
  bench_metric(
    "L2 scale",
    @moonhnsw.Distance::L2,
    10000,
    dim,
    50,
    k,
    m,
    ef_construction,
    20261027UL,
  )
  bench_metric(
    "L2 SQ8",
    @moonhnsw.Distance::L2,
    2000,
    dim,
    100,
    k,
    m,
    ef_construction,
    20261024UL,
    storage=@moonhnsw.Storage::SQ8,
  )
  bench_cosine_helper()
}

///|
fn main raise {
  run_all()
}