use super::*;

pub struct Markov {
    freq: HashMap<String, Vec<(String, usize)>>,
}

impl Markov {
    pub fn new(corpus: &str) -> Self {
        let mut freq = HashMap::<String, HashMap<String, usize>>::default();
        for (a, b) in corpus
            .split_whitespace()
            .map(|s| s.trim_matches(|c: char| !c.is_alphabetic()))
            .tuple_windows()
        {
            *freq
                .entry(a.to_string())
                .or_default()
                .entry(b.to_string())
                .or_default() += 1
        }
        Self {
            freq: freq
                .into_iter()
                .map(|(k, v)| (k, v.into_iter().collect::<Vec<_>>()))
                .collect::<HashMap<_, _>>(),
        }
    }
}

impl Generator for Markov {
    fn word_stream<'a>(
        &'a self,
        approx_len: usize,
        rng: &'a mut crate::Rng,
    ) -> Box<dyn Iterator<Item = Cow<'a, str>> + 'a> {
        let mut word = self.freq.keys().choose(rng).unwrap();
        Box::new(
            std::iter::from_fn(move || {
                word = self
                    .freq
                    .get(&*word)
                    .and_then(|rhs| rhs.choose_weighted(rng, |(_, n)| *n).ok())
                    .map(|(rhs, _)| rhs)
                    .unwrap_or(word);

                Some(word.clone().into())
            })
            .take(approx_len),
        )
    }
}
