///|
/// Distance metric used by the index.
///
/// Search results always report a *smaller-is-closer* distance:
///
/// - `L2`: squared Euclidean distance.
/// - `InnerProduct`: negated inner product (smaller = more similar).
/// - `Cosine`: negated inner product of L2-normalized vectors; vectors are
///   normalized once on insertion and queries are normalized per search.
pub(all) enum Distance {
  L2
  InnerProduct
  Cosine
} derive(Eq)

///|
pub extend Distance with Eq::{not_equal, equal}

///|
/// Tuning parameters of the index. Construct with `HnswParams::new` to get
/// defaults; values outside the supported range are clamped by `Hnsw::new`.
pub struct HnswParams {
  /// Maximum number of links kept per node on layers above 0. Larger values
  /// give higher recall at the cost of memory and build time. Clamped to >= 2;
  /// values above 1,073,741,823 are rejected to keep the layer-0 budget safe.
  m : Int32
  /// Size of the candidate list kept while inserting a node. Larger values
  /// build a higher-quality graph at a higher build cost. Clamped to >= 1.
  ef_construction : Int32
  /// Seed of the internal PRNG used to draw node levels.
  seed : UInt64
  /// Vector storage precision (full floats or 8-bit scalar quantization).
  storage : Storage
}

///|
/// Default parameters: `m = 16`, `ef_construction = 200`, full-precision
/// storage, fixed random seed.
pub fn HnswParams::new(
  m~ : Int32 = 16,
  ef_construction~ : Int32 = 200,
  seed~ : UInt64 = 0x9E3779B97F4A7C15UL,
  storage~ : Storage = Float32,
) -> HnswParams {
  { m, ef_construction, seed, storage, }
}

///|
/// Errors reported by the index.
pub suberror HnswError {
  /// Invalid argument, including a NaN or infinite vector component.
  InvalidParams(String)
  /// A vector dimension does not match the index dimension.
  DimMismatch(Int32, Int32)
  /// Serialized data is malformed.
  CorruptData(String)
  /// A vector id does not exist or refers to a deleted vector.
  InvalidId(Int32)
}

///|
/// `Show` renders the same human-readable text as `message`, so errors print
/// directly: `println("failed: \{err}")`.
pub impl Show for HnswError with fn output(self, logger) {
  logger.write_string(self.message())
}

///|
pub extend HnswError with Show::{to_string, output}

///|
/// Human-readable description of the error.
pub fn HnswError::message(self : HnswError) -> String {
  match self {
    InvalidParams(msg) => "invalid params: " + msg
    DimMismatch(expected, got) =>
      "dimension mismatch: index expects \{expected}, got \{got}"
    CorruptData(msg) => "corrupt data: " + msg
    InvalidId(id) => "invalid vector id \{id} (out of range or already deleted)"
  }
}

///|
/// One search result: the vector id and its distance to the query under the
/// index metric (see `Distance` for sign conventions).
pub struct Neighbor {
  id : Int32
  distance : Double
} derive(Eq, @debug.Debug, ToJson)

///|
pub extend Neighbor with Eq::{not_equal, equal}

///|
pub extend Neighbor with @moonbitlang/core/debug.Debug::{to_repr}

///|
/// `Show` prints a compact `Neighbor(id, distance)` form for quick logging.
pub impl Show for Neighbor with fn output(self, logger) {
  logger.write_string("Neighbor(")
  logger.write_object(self.id)
  logger.write_string(", ")
  logger.write_object(self.distance)
  logger.write_string(")")
}

///|
pub extend Neighbor with Show::{to_string, output}

///|
pub extend Neighbor with ToJson::{to_json}

///|
/// An approximate nearest-neighbor index over fixed-dimension `f32` vectors,
/// using the HNSW navigable small-world graph (Malkov & Yashunin, 2016).
///
/// The index supports incremental insertion and top-k search on all backends
/// (native / wasm / js). It is single-threaded: do not share one index across
/// concurrent writers or readers.
pub struct Hnsw {
  /// Dimension of the stored vectors.
  dim : Int32
  /// Distance metric.
  distance : Distance
  /// Effective parameters (after clamping).
  params : HnswParams
  /// True when the metric pre-normalizes vectors (Cosine).
  normalize : Bool
  /// 1 / ln(m), the level assignment factor.
  level_mult : Double
  /// Id of the entry point, -1 when the index is empty.
  mut entry : Int32
  /// Highest layer any node participates in, -1 when the index is empty.
  mut max_level : Int32
  /// Total number of vectors ever inserted (including deleted tombstones).
  mut count : Int32
  /// Number of non-deleted (searchable) vectors.
  mut alive_count : Int32
  /// Alive markers for tombstone deletion, keyed by node id.
  mut alive : FixedArray[Bool]
  /// User-facing label per node; defaults to the internal id.
  labels : Array[UInt64]
  /// Selected storage precision.
  storage : Storage
  /// Flattened full-precision storage: `count * dim` floats. Held until the
  /// SQ8 store freezes, then released.
  mut vectors : FixedArray[Float]
  /// Number of vector slots allocated in `vectors`.
  mut f_capacity : Int32
  /// Flattened quantized storage: `count * dim` bytes, valid once frozen.
  mut quants : FixedArray[Byte]
  /// Number of vector slots allocated in `quants`.
  mut q_capacity : Int32
  /// Per-dimension lower quantization bound (running until frozen).
  q_mins : FixedArray[Float]
  /// Per-dimension upper quantization bound (running until frozen).
  q_maxs : FixedArray[Float]
  /// Per-dimension scale, (max - min) / 255; valid once frozen.
  q_scales : FixedArray[Float]
  /// True once the quantization bounds are frozen and vectors are stored as
  /// one byte per dimension.
  mut q_frozen : Bool
  /// Reusable scratch for the normalized copy of an inserted vector.
  f_scratch : FixedArray[Float]
  /// Backing stores of the traversal heaps, drained after every use and
  /// reused so searches do not allocate fresh arrays.
  cand_items : Array[(Double, Int32)]
  result_items : Array[(Double, Int32)]
  /// `links[node][layer]` is the neighbor id list of `node` on `layer`.
  links : Array[Array[Array[Int32]]]
  /// Visited markers for graph traversal, keyed by node id.
  mut visited : FixedArray[Int32]
  /// Current traversal generation.
  mut stamp : Int32
  /// PRNG state used for level assignment.
  mut rng_state : UInt64
  /// Distance computations since the last search entry point: the
  /// hardware-independent efficiency metric of the graph.
  mut dist_calls : Int64
}

///|
/// Creates an empty index over vectors of dimension `dim` (1..=1,000,000).
pub fn Hnsw::new(
  dim : Int32,
  distance : Distance,
  params~ : HnswParams = HnswParams::new(),
) -> Hnsw raise HnswError {
  if dim < 1 {
    raise InvalidParams("dim must be >= 1")
  }
  if dim > 1000000 {
    raise InvalidParams("dim must be <= 1000000")
  }
  let m = if params.m < 2 { 2 } else { params.m }
  if m > 1073741823 {
    raise InvalidParams("m is too large")
  }
  let efc = if params.ef_construction < 1 { 1 } else { params.ef_construction }
  {
    dim,
    distance,
    params: HnswParams::new(
      m~,
      ef_construction=efc,
      seed=params.seed,
      storage=params.storage,
    ),
    normalize: match distance {
      Cosine => true
      _ => false
    },
    level_mult: 1.0 / @math.ln(m.to_double()),
    entry: -1,
    max_level: -1,
    count: 0,
    alive_count: 0,
    alive: FixedArray::make(0, false),
    labels: [],
    storage: params.storage,
    vectors: FixedArray::make(0, 0.0),
    f_capacity: 0,
    quants: FixedArray::make(0, 0),
    q_capacity: 0,
    q_mins: FixedArray::make(dim, Float::reinterpret_from_int(0x7F7FFFFF)),
    q_maxs: FixedArray::make(dim, -Float::reinterpret_from_int(0x7F7FFFFF)),
    q_scales: FixedArray::make(dim, 1.0),
    q_frozen: false,
    f_scratch: FixedArray::make(dim, 0.0),
    cand_items: [],
    result_items: [],
    links: [],
    visited: FixedArray::make(0, 0),
    stamp: 0,
    rng_state: params.seed,
    dist_calls: 0L,
  }
}

///|
/// Distance computations performed by the last search entry point: a
/// hardware-independent work count. Compare counts under the same dataset,
/// query set, metric, and search settings. Composing
/// entry points like `search_mmr` or `search_batch` report the total of
/// their inner searches plus their own re-ranking work.
pub fn Hnsw::distance_computations(self : Hnsw) -> Int64 {
  self.dist_calls
}

///|
/// Number of non-deleted vectors currently searchable in the index.
pub fn Hnsw::size(self : Hnsw) -> Int32 {
  self.alive_count
}

///|
/// Total number of vectors ever inserted, including deleted tombstones.
pub fn Hnsw::total(self : Hnsw) -> Int32 {
  self.count
}

///|
/// Returns true when `id` refers to a stored, non-deleted vector.
pub fn Hnsw::contains(self : Hnsw, id : Int32) -> Bool {
  id >= 0 && id < self.count && self.alive[id]
}

///|
/// Returns the label stored for `id`: the id itself by default, or the
/// value passed to `add_labeled`. Raises `InvalidId` for unknown or deleted
/// ids.
pub fn Hnsw::label_of(self : Hnsw, id : Int32) -> UInt64 raise HnswError {
  if !self.contains(id) {
    raise InvalidId(id)
  }
  self.labels[id]
}

///|
/// Returns a copy of the stored vector for `id`.
///
/// For `Cosine` indexes this is the normalized vector as used internally.
/// Raises `InvalidId` for unknown or deleted ids.
pub fn Hnsw::get_vector(
  self : Hnsw,
  id : Int32,
) -> FixedArray[Float] raise HnswError {
  if !self.contains(id) {
    raise InvalidId(id)
  }
  let out : FixedArray[Float] = FixedArray::make(self.dim, 0.0)
  if self.q_active() {
    self.dequantize_into(id, out)
  } else {
    FixedArray::unsafe_blit(out, 0, self.vectors, id * self.dim, self.dim)
  }
  out
}

///|
/// Diagnostic snapshot of the index graph.
pub struct HnswStats {
  /// Total insertions, including deleted tombstones.
  count : Int32
  /// Non-deleted (searchable) vectors.
  alive : Int32
  /// Highest layer in use; -1 for an empty index.
  max_level : Int32
  /// Number of nodes participating on each layer (index 0..=max_level).
  /// Deleted nodes still participate: their edges keep the graph connected.
  nodes_per_level : Array[Int32]
  /// Average number of links per participating node on each layer.
  avg_degree_per_level : Array[Double]
  /// Rough estimate of the heap memory held by the index, in bytes: vector
  /// storage, edge ids, per-list overhead and traversal markers.
  estimated_bytes : Int64
} derive(Eq, @debug.Debug, ToJson)

///|
pub extend HnswStats with Eq::{not_equal, equal}

///|
pub extend HnswStats with @moonbitlang/core/debug.Debug::{to_repr}

///|
pub extend HnswStats with ToJson::{to_json}

///|
/// Collects diagnostic statistics about the graph: how many nodes live on
/// each layer and how densely they are connected. Handy when tuning `m` and
/// `ef_construction`.
pub fn Hnsw::stats(self : Hnsw) -> HnswStats {
  let levels = if self.max_level < 0 { 0 } else { self.max_level + 1 }
  let nodes : Array[Int32] = Array::make(levels, 0)
  let degree_sum : Array[Int32] = Array::make(levels, 0)
  let mut lists = 0
  let mut edges = 0
  for node in self.links {
    let top = node.length() - 1
    let mut l = 0
    while l <= top {
      nodes[l] += 1
      degree_sum[l] += node[l].length()
      lists += 1
      edges += node[l].length()
      l += 1
    }
  }
  let avg : Array[Double] = []
  for l in 0.. Array[Int32] {
  let out : Array[Int32] = []
  for id in 0.. Hnsw raise HnswError {
  let idx = Hnsw::new(dim, distance, params~)
  for v in vectors {
    let _ = idx.add(v)
  }
  idx
}

///|
/// Same payload as `save`, wrapped as an immutable `Bytes` value for
/// convenient handling alongside other binary data.
pub fn Hnsw::save_bytes(self : Hnsw) -> Bytes {
  Bytes::from_array(self.save())
}

///|
/// Restores an index from a `Bytes` payload produced by `save_bytes`.
/// Raises `CorruptData` when the payload is truncated or inconsistent.
pub fn Hnsw::load_bytes(data : Bytes) -> Hnsw raise HnswError {
  Hnsw::load(data.to_fixedarray())
}

///|
/// Link budget of existing nodes on a layer: layer 0 keeps twice as many
/// links as upper layers.
fn Hnsw::max_links(self : Hnsw, layer : Int32) -> Int32 {
  if layer == 0 {
    self.params.m * 2
  } else {
    self.params.m
  }
}