extern crate naivebayes;

use naivebayes::NaiveBayes;

fn to_vec(line: &str) -> Vec<String> {
    let tokens = line.split(" ");
    let mut v: Vec<String> = Vec::new();
    for token in tokens {
        v.push(token.to_string());
    }
    return v;
}

fn main() {
    let mut nb = NaiveBayes::new();
    nb.train(&to_vec("great product"), &"positive".to_string());
    nb.train(
        &to_vec("the protection level is poor"),
        &"negative".to_string(),
    );
    nb.train(
        &to_vec("this is a great band I love them"),
        &"positive".to_string(),
    );
    nb.train(
        &to_vec("never buy this product it is too bad"),
        &"negative".to_string(),
    );
    nb.train(&to_vec("i love the shoes"), &"positive".to_string());
    nb.train(
        &to_vec("good product happy with the purchase"),
        &"positive".to_string(),
    );
    let to_classify: Vec<String> = to_vec("I love it is great");
    let classification = nb.classify(&to_classify);
    print!("classification = {:?}", classification);
}