///|
fn main {
  println("==================================================")
  println("    MoonVector Database & Index Demonstration     ")
  println("==================================================")
  println("Project Name : MoonVector Search Engine")
  println("Namespace    : ywz1314/vector")
  println("Description  : A zero-dependency vector similarity search")
  println("               and index database in MoonBit.")
  println("==================================================\n")

  // 1. Prepare synthetic dataset of 5-dimensional vectors representing document embeddings
  let docs = [
    @vector.Document::new("doc_1", [1.0, 2.0, 3.0, 4.0, 5.0], [
      ("author", "ywz1314"),
      ("category", "ml"),
    ]),
    @vector.Document::new("doc_2", [1.5, 2.5, 3.5, 4.5, 5.5], [
      ("author", "ywz1314"),
      ("category", "ml"),
    ]),
    @vector.Document::new("doc_3", [1.0, -1.0, 2.0, -2.0, 3.0], [
      ("author", "ywz1314"),
      ("category", "security"),
    ]),
    @vector.Document::new("doc_4", [-1.0, -2.0, -3.0, -4.0, -5.0], [
      ("author", "ywz1314"),
      ("category", "security"),
    ]),
    @vector.Document::new("doc_5", [0.0, 0.0, 0.0, 0.0, 1.0], [
      ("author", "ywz1314"),
      ("category", "general"),
    ]),
  ]

  println("Dataset loaded: 5 documents in 5-dimensional embedding space.\n")

  // Demo 1: FlatIndex with Cosine Similarity and Metadata Filter
  run_flat_demo(docs)

  // Demo 2: IVF-Flat (Inverted File) Cluster Search
  run_ivf_demo(docs)

  // Demo 3: KD-Tree Spatial Partition Nearest Neighbor Search
  run_kdtree_demo(docs)

  // Demo 4: LshIndex (Locality Sensitive Hashing) Approximate Cosine Search
  run_lsh_demo(docs)
}

///|
fn vector_to_string(v : Array[Double]) -> String {
  let sb = StringBuilder::new()
  sb.write_string("[")
  for i = 0; i < v.length(); i = i + 1 {
    if i > 0 {
      sb.write_string(", ")
    }
    sb.write_string(v[i].to_string())
  }
  sb.write_string("]")
  sb.to_string()
}

///|
fn run_flat_demo(docs : Array[@vector.Document]) -> Unit {
  println("--------------------------------------------------")
  println("    Demo 1: Flat baseline Index (Exact Search)    ")
  println("--------------------------------------------------")

  let index = @vector.FlatIndex::new()
  for doc in docs {
    index.add(doc)
  }

  println("Index Stats: " + index.stats().to_string())

  let query = [1.2, 2.2, 3.2, 4.2, 5.2]
  println("Query Vector : " + vector_to_string(query))

  try {
    // Search Top 3 by Cosine Similarity without filters
    println("\n[Search: Cosine Similarity, Top 3]")
    let res1 = index.search(query, 3, Cosine, [])
    print_results(res1)

    // Search Top 3 by Cosine Similarity with filter: category == "security"
    println(
      "\n[Search: Cosine Similarity, Top 3, Filter: category == \"security\"]",
    )
    let res2 = index.search(query, 3, Cosine, [("category", "security")])
    print_results(res2)

    // Search Top 3 by Euclidean Distance
    println("\n[Search: L2 Euclidean Distance, Top 3 (lower score is closer)]")
    let res3 = index.search(query, 3, Euclidean, [])
    print_results(res3)
  } catch {
    e => println("Flat search error: " + e.to_string())
  }
  println("")
}

///|
fn run_ivf_demo(docs : Array[@vector.Document]) -> Unit {
  println("--------------------------------------------------")
  println("    Demo 2: IVF-Flat Inverted File Cluster Index   ")
  println("--------------------------------------------------")

  try {
    let index = @vector.IvfIndex::new(2, Cosine)
    println("Building IVF index with K=2 centroids using K-Means...")
    index.build(docs)
    println("IVF index built successfully!")
    println("Index Stats: " + index.stats().to_string())

    let query = [1.2, 2.2, 3.2, 4.2, 5.2]
    println("Query Vector : " + vector_to_string(query))

    // Search with nprobe=1 (search only closest cluster centroid)
    println("\n[Search: IVF-Flat Cosine, Top 3, nprobe=1]")
    let res1 = index.search(query, 3, 1, [])
    print_results(res1)

    // Search with nprobe=2 (search both clusters - matches Flat baseline)
    println("\n[Search: IVF-Flat Cosine, Top 3, nprobe=2]")
    let res2 = index.search(query, 3, 2, [])
    print_results(res2)
  } catch {
    e => println("IVF build/search error: " + e.to_string())
  }
  println("")
}

///|
fn run_kdtree_demo(docs : Array[@vector.Document]) -> Unit {
  println("--------------------------------------------------")
  println("   Demo 3: KD-Tree Spatial Partition Index        ")
  println("--------------------------------------------------")

  try {
    let index = @vector.KdTreeIndex::new(Euclidean)
    println("Building KD-Tree index recursively...")

    index.build(docs)
    println("KD-Tree built successfully!")
    println("Index Stats: " + index.stats().to_string())

    let query = [1.2, 2.2, 3.2, 4.2, 5.2]
    println("Query Vector : " + vector_to_string(query))

    // Search Top 3 using KD-Tree exact nearest neighbor traversal
    println("\n[Search: KD-Tree L2 Distance, Top 3]")
    let res = index.search(query, 3, [])
    print_results(res)
  } catch {
    e => println("KD-Tree error: " + e.to_string())
  }
  println("")
}

///|
fn run_lsh_demo(docs : Array[@vector.Document]) -> Unit {
  println("--------------------------------------------------")
  println("   Demo 4: LSH Projection Hashing Index           ")
  println("--------------------------------------------------")

  try {
    let index = @vector.LshIndex::new(4, 5)
    println("Indexing documents into LSH projection buckets...")
    for doc in docs {
      index.add(doc)
    }
    println("LSH indexing completed!")

    let query = [1.2, 2.2, 3.2, 4.2, 5.2]
    println("Query Vector : " + vector_to_string(query))

    // Search Top 3 using LSH approximate search
    println("\n[Search: LSH Cosine approximation, Top 3]")
    let res = index.search(query, 3, [])
    print_results(res)
  } catch {
    e => println("LSH error: " + e.to_string())
  }
  println("==================================================")
}

///|
fn print_results(results : Array[@vector.SearchResult]) -> Unit {
  if results.length() == 0 {
    println("  No results matching filter requirements.")
    return
  }
  for i = 0; i < results.length(); i = i + 1 {
    let r = results[i]
    let mut cat = ""
    for m in r.metadata {
      if m.0 == "category" {
        cat = m.1
      }
    }
    println(
      "  " +
      (i + 1).to_string() +
      ". Doc ID: " +
      r.id +
      " (Score/Distance: " +
      r.score.to_string() +
      ", Category: " +
      cat +
      ")",
    )
  }
}