use std::collections::HashMap;

fn solve_part1(folts: &[usize]) -> (usize, usize) {
    let counts = folts
        .iter()
        .fold((0, HashMap::<usize, usize>::new()), |mut acc, elt| {
            match elt - acc.0 {
                x if (1..=3).contains(&x) => {
                    if let Some(count) = acc.1.get_mut(&x) {
                        *count += 1;
                    } else {
                        acc.1.insert(x, 1);
                    }
                    (*elt, acc.1)
                }
                _ => unreachable!()
            }
    }).1;
    // Counts of 3 should be incremented by one given last adapter's spec
    (*counts.get(&1).unwrap(), *counts.get(&3).unwrap() + 1)
}

fn solve_part2(folts: &[usize]) -> usize {
    let mut counts: Vec<usize> = std::iter::once(0usize)
        .chain(folts
                .iter().copied())
        .collect();
    counts.sort_unstable();
    counts.push(counts.iter().max().unwrap() + 3);
    // Map <folt, nb of moves from folt>
    let mut diffs: HashMap<usize, usize> = HashMap::new();
    diffs.insert(*counts.iter().max().unwrap(), 1);
    for &foltage in counts.iter().rev().skip(1) {
        let possible_steps = counts
            .iter()
            .filter(|&x| (1..=3).contains(&(*x as isize - foltage as isize)));
        diffs.insert(foltage, possible_steps.map(|x| diffs.get(x).unwrap()).sum());
    }
    // dbg!(&diffs);
    *diffs.get(&0).unwrap()
}

fn main() {
    let input = include_str!("../data/input");
    let mut folts: Vec<_> = input.lines().map(|l| l.parse::<usize>().unwrap()).collect();
    folts.sort();
    let counts = solve_part1(&folts);
    let counts2 = solve_part2(&folts);
    
    println!("Counts : {:?}, solution {}",
        counts,
        counts.0 * counts.1);
    println!("Number of paths possible : {}", counts2);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_step1() {
        let input = r"28
33
18
42
31
14
46
20
48
47
24
23
49
45
19
38
39
11
1
32
25
35
8
17
7
9
4
2
34
10
3";
        let mut folts: Vec<_> = input.lines().map(|l| l.parse::<usize>().unwrap()).collect();
        folts.sort();
        assert_eq!(solve_part1(&folts), (22, 10));
    }
}