use std::time::Instant;

use rand::{thread_rng, Rng};
use vectordb::Vector;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let db = vectordb::VectorDb::open("./vectordb", 2)?;

    let total_chunks = 32;
    let chunk_size = 1024 * 8;

    for _ in 0..total_chunks {
        let vectors = (0..chunk_size)
            .map(|_| thread_rng().gen::<[f32; 7]>().into())
            .collect::<Vec<_>>();

        println!("Inserting vectors");
        let now = Instant::now();
        db.insert_vectors_blocking(vectors)?;
        println!("Inserting vectors - Took {:?}", now.elapsed());
        std::thread::sleep(std::time::Duration::from_secs(3));
    }

    println!();

    std::thread::sleep(std::time::Duration::from_secs(10));

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

    db.insert_vector_blocking(existing_vector.clone())?;

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

    for v in &test_vecs {
        println!("Finding similarities");
        let now = Instant::now();
        let similarities = db.find_similarities_blocking(v.clone(), None, 5)?;
        println!("Finding similarities - Took {:?}", now.elapsed());

        for id in similarities {
            let similar = db.get_vector_blocking(id)?.expect("Vector exists");
            println!(
                "similar to {}: {}",
                id,
                v.squared_euclidean_distance(&similar)
            );
        }
    }

    println!();

    for v in &test_vecs {
        println!("Finding furthest similarities");
        let now = Instant::now();
        let similarities = db.find_furthest_similarities_blocking(v.clone(), None, 5)?;
        println!("Finding furthest similarities - Took {:?}", now.elapsed());

        for id in similarities {
            let similar = db.get_vector_blocking(id)?.expect("Vector exists");
            println!(
                "similar to {}: {}",
                id,
                v.squared_euclidean_distance(&similar)
            );
        }
    }

    println!();

    for v in &test_vecs {
        println!("Finding dissimilarities");
        let now = Instant::now();
        let similarities = db.find_dissimilarities_blocking(v.clone(), None, 5)?;
        println!("Finding dissimilarities - Took {:?}", now.elapsed());

        for id in similarities {
            let similar = db.get_vector_blocking(id)?.expect("Vector exists");
            println!(
                "dissimilar to {}: {}",
                id,
                v.squared_euclidean_distance(&similar)
            );
        }
    }

    println!();

    for v in &test_vecs {
        println!("Finding closest dissimilarities");
        let now = Instant::now();
        let similarities = db.find_closest_dissimilarities_blocking(v.clone(), None, 5)?;
        println!("Finding closest dissimilarities - Took {:?}", now.elapsed());

        for id in similarities {
            let similar = db.get_vector_blocking(id)?.expect("Vector exists");
            println!(
                "dissimilar to {}: {}",
                id,
                v.squared_euclidean_distance(&similar)
            );
        }
    }

    Ok(())
}