use std::time::Instant;

use itertools::Itertools;
use rand::{thread_rng, Rng};

fn main() {
    let db = vectordb::VectorDb::open("./db.redb", String::from("moments")).expect("Launched");

    let total = 1024 * 1024;

    println!("Inserting");
    let now = Instant::now();

    let chunks = (0..total)
        .map(|_| thread_rng().gen::<[f32; 7]>().into())
        .chunks(1024);

    let iter = chunks.into_iter().map(|chunk| {
        let vectors = chunk.collect::<Vec<_>>();
        db.insert_many_vectors(&vectors)
    });

    for res in iter {
        res.expect("Insert works");
    }
    println!("Inserting - Took {:?}", now.elapsed());

    let existing_vector = thread_rng().gen::<[f32; 7]>().into();

    db.insert_vector(&existing_vector).expect("Insert works");

    for _ in 0..4 {
        println!("Creating Hypertree");
        let now = Instant::now();
        db.create_hypertree().expect("Tree created");
        println!("Creating Hypertree - Took {:?}", now.elapsed());
    }

    let test_vecs = [existing_vector]
        .into_iter()
        .chain((0..50).map(|_| thread_rng().gen::<[f32; 7]>().into()));

    for v in test_vecs {
        let compare_id = db.insert_vector(&v).expect("Inserted");

        println!("Finding similarities");
        let now = Instant::now();
        let similarities = db
            .find_similarities(compare_id, 5)
            .expect("Similar vectors exist");
        println!("Finding similarities - Took {:?}", now.elapsed());

        for id in similarities {
            let similar = db
                .get_vector(id)
                .expect("Got vector")
                .expect("Vector exists");
            println!(
                "{} is similar to {}: {}",
                compare_id,
                id,
                v.squared_euclidean_distance(&similar)
            );
        }
    }
}