use aoc_runner_derive::{aoc, aoc_generator};
use std::collections::{HashMap, HashSet};
#[derive(Debug, Eq, PartialEq, Clone)]
struct SearchState {
    available: HashSet<u16>,
    node: u16,
}

fn solve_part_2_rec(
    node: u16,
    goal: u16,
    map: &HashMap<u16, HashSet<u16>>,
    used: &mut HashSet<u16>,
    cache: &mut HashMap<(u16, u16), usize>,
) -> usize {
    if node == goal {
        return 1;
    }
    let mut ret = 0;
    // println!("+{}", node);
    for next in map.get(&node).iter().map(|v| v.iter()).flatten() {
        if used.contains(next) {
            continue;
        };
        used.insert(*next);
        let value = match cache.get(&(node, *next)) {
            Some(value) => *value,
            None => {
                let value = solve_part_2_rec(*next, goal, map, used, cache);
                cache.insert((node, *next), value);
                value
            }
        };
        ret += value;
        used.remove(next);
    }
    // println!("-{}", node);
    ret
}
pub struct Data(pub Vec<u16>);

impl Data {
    #[inline(always)]
    fn count_paths(&self) -> usize {
        let mut nodes = self.0.clone();
        nodes.push(0);
        nodes.push(nodes.iter().max().unwrap() + 3);
        nodes.sort_unstable();
        let goal = *nodes.last().unwrap();
        let mut map: HashMap<u16, HashSet<u16>> = HashMap::new();
        for n in &nodes {
            for v in &nodes {
                let d = v - n;
                if d > 0 && d < 4 {
                    map.entry(*n).or_default().insert(*v);
                }
            }
        }
        return solve_part_2_rec(0, goal, &map, &mut HashSet::new(), &mut HashMap::new());
    }
}

#[aoc_generator(day10)]
pub fn input_generator(input: &str) -> Data {
    let mut numbers = input
        .lines()
        .map(|l| l.trim().parse())
        .collect::<Result<Vec<u16>, _>>()
        .unwrap();
    numbers.sort_unstable();
    return Data(numbers);
}

#[aoc(day10, part1)]
#[inline(always)]
pub fn solve_part1(input: &Data) -> usize {
    let mut n_1 = 0;
    let mut n_3 = 0;
    let mut nums = input.0.clone();
    nums.push(0);
    nums.push(nums.iter().max().unwrap() + 3);
    nums.sort_unstable();
    for w in nums.windows(2) {
        match w[1] - w[0] {
            1 => n_1 += 1,
            3 => n_3 += 1,
            _ => (),
        }
    }
    return n_1 * n_3;
}

#[aoc(day10, part2)]
#[inline(always)]
pub fn solve_part2(input: &Data) -> usize {
    return input.count_paths();
}