markov2.rs

use std::io;
use std::io::BufferedReader;
use std::collections::HashMap;
use std::io::File;
use std::rand::{task_rng, Rng};

struct MarkovModels<'a> {
    data: HashMap<(&'a str, &'a str), Vec<&'a str>>,
    words: &'a [&'a str]
}

impl<'a> MarkovModels<'a> {
    pub fn new(words: &'a [&'a str]) -> MarkovModels<'a> {
        let mut hm : HashMap<(&str, &str), Vec<&str>> = HashMap::new();
        for i in range(0u, words.len() - 2) {
            let key = (*words.get(i).unwrap(), *words.get(i+1).unwrap());
            if !hm.contains_key(&key) {
                hm.insert(key, Vec::new());
            }
            let mut v : &mut Vec<&str> = hm.get_mut(&key).unwrap();
            v.push(*words.get(i+2).unwrap());
        }
        MarkovModels{
            data: hm,
            words: words
        }
    }

    pub fn tell_tail(&self, size: uint) -> Vec<&'a str> {
        let mut res = Vec::new();
        let mut task = task_rng();
        let starter = task.gen_range(0u, self.words.len());

        let mut first : &str = *self.words.get(starter).unwrap();
        let mut second : &str = *self.words.get(starter + 1).unwrap();
        for _ in range(0u, size) {
            res.push(first);

            let key = (first, second);
            first = second;
            second = {
                let value = self.data.get(&key).unwrap();
                *task.choose(value.as_slice()).unwrap()
            };
        }

        return res;
    }
}

fn main() {
    let path = Path::new("pg100.txt");
    let mut file = BufferedReader::new(File::open(&path));

    let mut data : Vec<String> = Vec::new();
    for line in file.lines() {
        let unwrapped_line : String = line.unwrap();
        for word in unwrapped_line.split_str(" ") {
            let cloned_trimmed_word : String = word.trim_right().trim_left().to_string();
            if cloned_trimmed_word.len() > 0 {
                data.push(cloned_trimmed_word);
            }
        }
    }

    let mapped_data : Vec<&str> = data.iter().map(|el| {
        el.as_slice()
    }).collect();
    let mm = MarkovModels::new(mapped_data.as_slice());

    let mut reader = io::stdin();
    loop {
        let result : String = mm.tell_tail(50).iter().fold("".to_string(), |a, b| {
            a + " " + b
        });
        println!("{}", result);
        reader.read_byte();
    }
}