// Locality Sensitive Hashing (LSH) Index for approximate Cosine Similarity search

///|
pub struct LshIndex {
  planes : Array[Array[Double]] // random projection planes
  buckets : Map[Int, Array[Document]] // hash code -> documents bucket
  num_planes : Int
  dim : Int
}

///|
pub fn LshIndex::new(num_planes : Int, dim : Int) -> LshIndex {
  let planes = Array::make(num_planes, [])
  for i = 0; i < num_planes; i = i + 1 {
    let plane = Array::make(dim, 0.0)
    // Fill planes deterministically for reproducible behavior
    for j = 0; j < dim; j = j + 1 {
      plane[j] = if (i + j) % 2 == 0 { 1.0 } else { -1.0 }
    }
    planes[i] = plane
  }
  { planes, buckets: {}, num_planes, dim }
}

///|
fn LshIndex::compute_hash(
  self : LshIndex,
  v : Array[Double],
) -> Int raise VectorError {
  let mut hash = 0
  for i = 0; i < self.num_planes; i = i + 1 {
    let dp = dot_product(v, self.planes[i])
    if dp >= 0.0 {
      hash = hash | (1 << i)
    }
  }
  hash
}

///|
pub fn LshIndex::add(self : LshIndex, doc : Document) -> Unit raise VectorError {
  if doc.vector.length() != self.dim {
    raise DimensionMismatch("Vector dimension mismatch in LSH index")
  }
  let hash = self.compute_hash(doc.vector)
  let bucket = match self.buckets.get(hash) {
    Some(b) => b
    None => {
      let b = []
      self.buckets.set(hash, b)
      b
    }
  }
  bucket.push(doc)
}

///|
pub fn LshIndex::search(
  self : LshIndex,
  query : Array[Double],
  top_k : Int,
  filters : Array[(String, String)],
) -> Array[SearchResult] raise VectorError {
  if query.length() != self.dim {
    raise DimensionMismatch("Query dimension mismatch in LSH search")
  }

  let q_hash = self.compute_hash(query)
  let candidates = []

  // Scan buckets within Hamming distance <= 1
  for entry in self.buckets {
    let hash = entry.0
    let docs_in_bucket = entry.1

    let diff = hash ^ q_hash
    let mut hamming = 0
    for i = 0; i < self.num_planes; i = i + 1 {
      if ((diff >> i) & 1) == 1 {
        hamming = hamming + 1
      }
    }

    if hamming <= 1 {
      for doc in docs_in_bucket {
        candidates.push(doc)
      }
    }
  }

  // Deduplicate and score candidates
  let results = []
  let seen = Map([])
  for doc in candidates {
    if !seen.contains(doc.id) {
      seen.set(doc.id, true)
      if matches_filters(doc.metadata, filters) {
        let score = cosine_similarity(query, doc.vector)
        results.push({ id: doc.id, score, metadata: doc.metadata })
      }
    }
  }

  // Sort descending for Cosine similarity
  sort_results(results, false)

  let limit = if top_k < results.length() { top_k } else { results.length() }
  let top = []
  for i = 0; i < limit; i = i + 1 {
    top.push(results[i])
  }
  top
}