use std::cell::Cell;
use std::collections::{HashMap, HashSet};

fn parse(s: &str) -> impl Iterator<Item=(&str, &str)> {
	s.split_ascii_whitespace().map(|l| {
		let mut x = l.splitn(2, ')');
		let a = x.next().unwrap();
		let b = x.next().unwrap();
		(a, b)
	})
}

fn count_orbits_get(orbit_relations: &HashMap<&str, (&str, Cell<Option<u32>>)>, object: &str) -> u32 {
	if object == "COM" { return 0; }
	let around = &orbit_relations[object];
	if let Some(n) = around.1.get() { return n; }
	let n = count_orbits_get(orbit_relations, around.0) + 1;
	around.1.set(Some(n));
	n
}

fn count_orbits(input: &str) -> u32 {
	let orbit_relations: HashMap<&str, (&str, Cell<Option<u32>>)> = parse(input).map(|(a, b)| (b, (a, Cell::new(None)))).collect();
	let mut sum = 0;
	for (_object, around) in &orbit_relations {
		if let Some(n) = around.1.get() {
			sum += n;
		} else {
			let n = count_orbits_get(&orbit_relations, around.0) + 1;
			around.1.set(Some(n));
			sum += n;
		}
	}
	// println!("Map: {:?}", orbit_relations);
	sum
}

fn count_orbit_transfers(input: &str) -> u32 {
	let orbit_relations: HashMap<&str, (&str, Cell<Option<u32>>)> = parse(input).map(|(a, b)| (b, (a, Cell::new(None)))).collect();
	for (_object, around) in &orbit_relations {
		if around.1.get().is_none() {
			let n = count_orbits_get(&orbit_relations, around.0) + 1;
			around.1.set(Some(n));
		}
	}
	// println!("Map: {:?}", orbit_relations);
	let mut you = HashSet::new();
	let mut object = "YOU";
	loop {
		if "COM" == object { break; }
		object = orbit_relations[object].0;
		// println!("YOU orbit (...) {}", object);
		you.insert(object);
	}
	object = "SAN";
	loop {
		if you.contains(&object) {
			let you = orbit_relations["YOU"].1.get().unwrap();
			let san = orbit_relations["SAN"].1.get().unwrap();
			let bottom = orbit_relations[object].1.get().unwrap();
			return (you - bottom - 1) + (san - bottom - 1);
		}
		object = orbit_relations[object].0;
		// println!("SAN orbit (...) {}", object);
	}
}

fn main() {
	let input = include_str!("input.txt");
	println!("Orbit sum: {}", count_orbits(input));
	println!("Orbit transfers: {}", count_orbit_transfers(input));
}

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

	#[test]
	fn examples1() {
		assert_eq!(count_orbits("COM)B B)C C)D D)E E)F B)G G)H D)I E)J J)K K)L"), 42);
	}

	#[test]
	fn examples2() {
		assert_eq!(count_orbit_transfers("COM)B B)C C)D D)E E)F B)G G)H D)I E)J J)K K)L K)YOU I)SAN"), 4);
	}
}