use rand::rngs::StdRng;
use rand::seq::SliceRandom;
use rand::{Rng, SeedableRng};

fn generate_boxes(n: usize, mut rng: impl Rng) -> Vec<usize> {
    let mut range: Vec<_> = (0..n).collect();
    range.shuffle(&mut rng);
    range
}

fn search_for_self(self_num: usize, limit: usize, boxes: &[usize]) -> bool {
    let mut cur_box = self_num;
    let mut num_opened = 0;
    while num_opened < limit {
        let num = boxes[cur_box];
        if num == self_num {
            return true;
        }
        cur_box = num;
        num_opened += 1;
    }
    false
}

fn sample(n: usize, boxes: &[usize]) -> bool {
    (0..n).all(|person| search_for_self(person, n / 2, boxes))
}

fn main() {
    const TRIES_LEN: usize = 45;
    const BOX_LEN: usize = 90;

    let seed = 12;
    let mut rng = StdRng::seed_from_u64(seed);
    let num_samples = 100;

    let max_tries_options: [usize; TRIES_LEN] = [
        5,  6,  7,  8,  9,  10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
        20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34,
        35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49
    ];

    let num_box_options: [usize; BOX_LEN] = [
        10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24,
        25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
        40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54,
        55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69,
        70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84,
        85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99
    ];

    let mut results = [[0; BOX_LEN]; TRIES_LEN];
    for &max_tries in max_tries_options.iter() {
        for &num_boxes in num_box_options.iter() {
            let successes = (0..num_samples)
                .map(|_| {
                    let boxes = generate_boxes(num_boxes, &mut rng);
                    sample(num_boxes, &boxes) as u32
                })
                .sum();
            results[max_tries - max_tries_options[0]][num_boxes - num_box_options[0]] = successes;
        }
    }

    println!("[");
    for row in results.iter() {
        println!("  {:?}", &row[..]);
    }
    println!("]");
}