use aoc_runner_derive::{aoc, aoc_generator};
use rustc_data_structures::fx::{FxHashMap, FxHashSet};

type Edges = FxHashMap<String, Vec<String>>;

#[aoc_generator(day11)]
fn parse(input: &str) -> Edges {
    let mut edges: Edges = FxHashMap::default();
    for l in input.lines() {
        let (k, rest) = l.split_once(": ").unwrap();
        for v in rest.split_ascii_whitespace() {
            edges.entry(k.to_string()).or_default().push(v.to_string());
        }
    }

    edges
}

fn find_path(cur: &str, goal: &str, edges: &Edges) -> u64 {
    if cur == goal {
        return 1;
    }
    if let Some(nexts) = edges.get(cur) {
        nexts.iter().map(|n| find_path(n, goal, edges)).sum()
    } else {
        0
    }
}

fn mark_paths<'a>(
    cur: &'a str,
    edges: &'a Edges,
    mut reachable: FxHashSet<&'a str>,
) -> FxHashSet<&'a str> {
    reachable.insert(cur);
    if let Some(nexts) = edges.get(cur) {
        for n in nexts {
            if !reachable.contains(&n.as_str()) {
                reachable = mark_paths(n, edges, reachable);
            }
        }
    }
    reachable
}

fn count_paths_impl<'a>(
    cur: &'a str,
    goal: &str,
    edges: &'a Edges,
    memo: &mut FxHashMap<&'a str, u64>,
) -> u64 {
    if cur == goal {
        return 1;
    }
    if let Some(v) = memo.get(cur) {
        return *v;
    }

    if let Some(nexts) = edges.get(cur) {
        let ret = nexts
            .iter()
            .map(|n| count_paths_impl(n, goal, edges, memo))
            .sum();
        memo.insert(cur, ret);
        ret
    } else {
        memo.insert(cur, 0);
        0
    }
}

fn count_paths(cur: &str, goal: &str, edges: &Edges) -> u64 {
    let mut cache = FxHashMap::default();
    count_paths_impl(cur, goal, edges, &mut cache)
}

#[aoc(day11, part1, NaiveDfs)]
fn part1(fwd: &Edges) -> u64 {
    find_path("you", "out", fwd)
}

#[aoc(day11, part1, MemoDfs)]
fn part1_memo(fwd: &Edges) -> u64 {
    count_paths("you", "out", fwd)
}

#[aoc(day11, part2, MemoDfs)]
fn part2_memo(edges: &Edges) -> u64 {
    // we assume all valid paths reach svr->fft->dac->out in that order. this holds on the example and the input.
    // to handle both possible orderings, we need to sum paths fft->dac and dac->fft, and then partition the graph
    // (just remove the node) at dac & fft before computing the sum of dac->out and fft->out.
    //
    // this is our performance implementation though and it's faster to not do that work. we do the posterity for the
    // naive solution
    count_paths("svr", "fft", edges)
        * count_paths("fft", "dac", edges)
        * count_paths("dac", "out", edges)
}

#[aoc(day11, part2, NaiveDfs)]
fn part2(edges: &Edges) -> u64 {
    let mut rev = Edges::from_iter(edges.keys().map(|k| (k.to_owned(), vec![])));
    for (from, tos) in edges {
        for to in tos {
            rev.entry(to.to_owned()).or_default().push(from.to_owned());
        }
    }
    let mut reachable_dac = mark_paths("dac", edges, FxHashSet::default());
    reachable_dac = mark_paths("dac", &rev, reachable_dac);
    let mut reachable_fft = mark_paths("fft", edges, FxHashSet::default());
    reachable_fft = mark_paths("fft", &rev, reachable_fft);

    let reachable: FxHashSet<&str> = reachable_dac
        .intersection(&reachable_fft)
        .copied()
        .collect();

    let unreachable: FxHashSet<&str> = FxHashSet::from_iter(edges.keys().map(|k| k.as_str()))
        .difference(&reachable)
        .copied()
        .collect();

    let mut reachable_edges = edges.clone();

    for k in &unreachable {
        reachable_edges.remove(*k);
    }
    for (_k, v) in reachable_edges.iter_mut() {
        for ur in &unreachable {
            if let Some(idx) = v.iter().position(|s| s == ur) {
                v.remove(idx);
            }
        }
    }

    // A real graph structure would probably notice that the graphs get split here
    let mut reachable_edges_excl_fft = reachable_edges.clone();
    reachable_edges_excl_fft.remove("fft");
    let mut reachable_edges_excl_dac = reachable_edges.clone();
    reachable_edges_excl_dac.remove("dac");

    find_path("svr", "fft", &reachable_edges)
        // only svr->fft->dac->out paths appear in my input and in the example, but include the dac->fft ordering
        // as well, for posterity. It ~doubles runtime.
        * (find_path("fft", "dac", &reachable_edges) + find_path("dac", "fft", &reachable_edges))
        * (find_path("dac", "out", &reachable_edges_excl_fft) + find_path("fft", "out", &reachable_edges_excl_dac))
}

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

    use super::*;

    const EXAMPLE: &str = "aaa: you hhh
you: bbb ccc
bbb: ddd eee
ccc: ddd eee fff
ddd: ggg
eee: out
fff: out
ggg: out
hhh: ccc fff iii
iii: out";

    const EXAMPLE2: &str = "svr: aaa bbb
aaa: fft
fft: ccc
bbb: tty
tty: ccc
ccc: ddd eee
ddd: hub
hub: fff
eee: dac
dac: fff
fff: ggg hhh
ggg: out
hhh: out";

    #[rstest]
    #[case(EXAMPLE, 5)]
    fn part1_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part1(&parse(input)), expected);
    }

    #[rstest]
    #[case(EXAMPLE, 5)]
    fn part1_memo_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part1_memo(&parse(input)), expected);
    }

    #[rstest]
    #[case(EXAMPLE2, 2)]
    fn part2_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part2(&parse(input)), expected);
    }

    #[rstest]
    #[case(EXAMPLE2, 2)]
    fn part2_memo_example(#[case] input: &str, #[case] expected: u64) {
        assert_eq!(part2_memo(&parse(input)), expected);
    }
}