use criterion::criterion_main;
use criterion::*;
use data_value::DataValue;
use ndarray::Array1;
use rand::*;
use trs_dataframe::{dataframe::TopN, df, DataFrame, DataType, JoinRelation, Key};

fn sort_copy(df: DataFrame, key: &Key) {
    let r = df.sorted(key);
    assert!(r.is_ok(), "{r:?}");
    let r = r.unwrap().get_sorted();
    assert_eq!(r.len(), df.len());
}

fn sort_only(df: DataFrame, key: &Key) {
    let r = df.sorted(key);
    assert!(r.is_ok(), "{r:?}");
}

fn join(mut df1: DataFrame, df2: DataFrame, join_releation: &JoinRelation) {
    let r = df1.join(df2, join_releation);
    assert!(r.is_ok(), "{r:?}");
}

fn topn(df: DataFrame, key: &Key, top: TopN) {
    let r = df.sorted(key);
    assert!(r.is_ok(), "{r:?}");
    let r = r.unwrap().topn(top);
    assert!(r.is_ok(), "{r:?}");
}

pub fn criterion_benchmark(c: &mut Criterion) {
    let mut df = DataFrame::default();
    let mut rng = rand::rng();
    for key in 0..20 {
        let v: Vec<DataValue> = (0..10_000)
            .map(|_x| rng.random::<i32>().into())
            .collect::<Vec<_>>();
        let _ = df.add_single_column(
            Key::new(&format!("a{key}"), DataType::U32),
            Array1::from_vec(v),
        );
    }
    let mut df2 = DataFrame::default();
    for key in 0..10 {
        let v: Vec<DataValue> = (0..10)
            .map(|_x| rng.random::<i32>().into())
            .collect::<Vec<_>>();
        let _ = df2.add_single_column(
            Key::new(&format!("a{key}"), DataType::U32),
            Array1::from_vec(v),
        );
    }

    let df3 = df![
        "a" => vec![1i64],
        "b" => vec![2i64],
        "c" => vec![3i64]
    ];

    let single = df![
        "z" => vec![1i64]
    ];

    let key = Key::new("a1", DataType::U32);
    c.bench_function("topn_first", |b| {
        b.iter(|| topn(df.clone(), &key, TopN::First(1000)))
    });
    c.bench_function("topn_last", |b| {
        b.iter(|| topn(df.clone(), &key, TopN::Last(1000)))
    });
    c.bench_function("sort", |b| b.iter(|| sort_copy(df.clone(), &key)));
    c.bench_function("sort_only", |b| b.iter(|| sort_only(df.clone(), &key)));
    c.bench_function("join_cartesian", |b| {
        b.iter(|| join(df.clone(), df2.clone(), &JoinRelation::cartesian_product()))
    });
    c.bench_function("join_cartesian_right", |b| {
        b.iter(|| join(df.clone(), df3.clone(), &JoinRelation::cartesian_product()))
    });
    c.bench_function("join_cartesian_left", |b| {
        b.iter(|| join(df3.clone(), df.clone(), &JoinRelation::cartesian_product()))
    });
    c.bench_function("join_broadcast", |b| {
        b.iter(|| join(df.clone(), single.clone(), &JoinRelation::broadcast()))
    });
    c.bench_function("join_extend", |b| {
        b.iter(|| join(df.clone(), df.clone(), &JoinRelation::extend()))
    });
    let mut dfl = DataFrame::default();
    for key in 0..10 {
        let v: Vec<DataValue> = (0..10)
            .map(|_x| rng.random::<i32>().into())
            .collect::<Vec<_>>();
        let _ = dfl.add_single_column(
            Key::new(&format!("a{key}"), DataType::U32),
            Array1::from_vec(v),
        );
    }
    let mut dfr = DataFrame::default();
    for key in 0..3 {
        let v: Vec<DataValue> = (0..10)
            .map(|_x| rng.random::<i32>().into())
            .collect::<Vec<_>>();
        let _ = dfr.add_single_column(
            Key::new(&format!("a{key}"), DataType::U32),
            Array1::from_vec(v),
        );
    }
    for key in 0..3 {
        let v: Vec<DataValue> = (0..10)
            .map(|_x| rng.random::<i32>().into())
            .collect::<Vec<_>>();
        let _ = dfr.add_single_column(
            Key::new(&format!("b{key}"), DataType::U32),
            Array1::from_vec(v),
        );
    }

    c.bench_function("join_by_id", |b| {
        b.iter(|| {
            join(
                dfl.clone(),
                dfr.clone(),
                &JoinRelation::join_by_id(vec!["a0".into(), "a1".into(), "a2".into()]),
            )
        });
    });
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
