A simple markov chain program written in rust

markov_chain_rust.rs

use rand::distributions::WeightedIndex;
use rand::prelude::*;

struct MarkovChain{
    transition_matrix: Vec<Vec<f64>>,
    states: Vec<String>,
}

impl MarkovChain {
    fn new(transition_matrix: Vec<Vec<f64>>, states: Vec<String>) -> MarkovChain {
        MarkovChain{
            transition_matrix,
            states,
        }
    }

    fn next_states(&self,current_state:&str) -> String {
        let current_state_index = self.states.iter().position(|s| s == current_state).unwrap();
        //println!("{}",current_state_index);
        let probablities = &self.transition_matrix[current_state_index];
        // for &probability in probablities {
        //     println!("{}", probability);
        // }        
        let dist = WeightedIndex::new(probablities).unwrap();
        let mut rng = thread_rng();
        let random_state_index = dist.sample(&mut rng);
        self.states[random_state_index].clone()
    }

    fn n_states(&self,starting_state:&str,n_states: usize) -> Vec<String> {

        let mut current_state = starting_state.to_string();
        let mut next_states_arr = Vec::new();

        for _ in 0..n_states {
            let next = self.next_states(&current_state);
            next_states_arr.push(next.clone());
            current_state = next.clone();
        }
        
        next_states_arr
    }
}

fn main(){
    let transition_matrix = vec![
        vec![0.2, 0.3, 0.5],
        vec![0.1, 0.4, 0.5],
        vec![0.3, 0.2, 0.5],
    ];

    let states = vec!["A".to_string(), "B".to_string(), "C".to_string()];

    let mc = MarkovChain::new(transition_matrix,states);
    let next = mc.next_states("B");
    println!("{}",next);
    let n_states_value = mc.n_states("A",5);
    for st in n_states_value {
        println!("{}",st);
    }
}