use ndarray::{Array2, Axis};
use rand::distr::Distribution;
use rand::distr::weighted::WeightedIndex;
use rand::prelude::*;
pub struct Markov {
pub matrix: Array2<f64>,
}
impl Markov {
pub fn train(names: &[&str]) -> Self {
let n = 256;
let mut matrix = Array2::<f64>::zeros((n, n));
for name in names {
let mut prev = b'^' as usize;
for c in name.chars() {
let j = c as usize;
if j < n {
matrix[[prev, j]] += 1.0;
prev = j;
}
}
let j = b'$' as usize; matrix[[prev, j]] += 1.0;
}
for mut row in matrix.axis_iter_mut(Axis(0)) {
let sum: f64 = row.sum();
if sum > 0.0 {
row /= sum;
}
}
Self { matrix }
}
pub fn precompute_distributions(&self) -> Vec<Option<WeightedIndex<f64>>> {
self.matrix
.axis_iter(Axis(0))
.map(|row| WeightedIndex::new(row.to_vec()).ok())
.collect()
}
pub fn generate(
&self,
rng: &mut impl Rng,
distributions: &Vec<Option<WeightedIndex<f64>>>,
) -> String {
let mut result = String::new();
let mut current = b'^' as usize;
loop {
let next = if let Some(dist) = &distributions[current] {
dist.sample(rng)
} else {
b'$' as usize
};
if next == b'$' as usize {
break;
}
result.push(next as u8 as char);
current = next;
}
result
}
}