use std::collections::HashMap;
use crate::Day;

type ReplacementMap = HashMap<String, Vec<String>>;

pub struct Day14 {
    start: HashMap<String, u64>,
    end: char,
    replacements: ReplacementMap,
}

impl Day for Day14 {
    fn init(content: String) -> anyhow::Result<Self> {
        let mut line_iter = content.lines().peekable();
        let input: Vec<char> = line_iter.next().unwrap().chars().collect();
        let mut start = HashMap::new();
        for i in 0..input.len() - 1 {
            let pair = format!("{}{}", input[i], input[i + 1]);
            let count = start.entry(pair).or_insert(0);
            *count += 1;
        }

        line_iter.next();

        let mut replacements = HashMap::new();
        for line in line_iter {
            let cap = line.split(" -> ").collect::<Vec<_>>();

            let pair = cap[0].to_string();
            let output = cap[1].to_string().chars().next().unwrap();
            replacements.insert(
                pair.clone(),
                vec![
                    format!("{}{}", pair.chars().next().unwrap(), output),
                    format!("{}{}", output, pair.chars().nth(1).unwrap()),
                ],
            );
        }

        Ok(Self {
            start,
            replacements,
            end: *input.last().unwrap(),
        })
    }

    fn part1(&self) -> anyhow::Result<String> {
        Ok(format!("{}", self.iterate(10)))
    }

    fn part2(&self) -> anyhow::Result<String> {
        Ok(format!("{}", self.iterate(40)))
    }
}

impl Day14 {
    fn iterate(&self, amount: u64) -> u64 {
        let mut s = self.start.clone();
        for _ in 0..amount {
            s = do_pain(&s, &self.replacements);
        }

        let mut counts: HashMap<char, u64> = HashMap::new();
        for (st, b) in s {
            {
                let e = counts.entry(st.chars().next().unwrap()).or_insert(0);
                *e += b;
            }
        }
        {
            let e = counts.entry(self.end).or_insert(0);
            *e += 1;
        }
        let mut min = u64::MAX;
        let mut max = 0;
        for &count in counts.values() {
            min = min.min(count);
            max = max.max(count);
        }
        max - min
    }
}

fn do_pain(input: &HashMap<String, u64>, map: &ReplacementMap) -> HashMap<String, u64> {
    let mut output_map = HashMap::new();
    for (pair, &count) in input {
        let new_pairs = map.get(pair).unwrap();

        for new_pair in new_pairs {
            let output_count = output_map.entry(new_pair.to_string()).or_insert(0);
            *output_count += count;
        }
    }

    output_map
}

#[cfg(test)]
mod tests {
    use crate::day14::Day14;
    use crate::day::Day;
    use crate::day_tests;

    const INPUT: &str = r"NNCB

CH -> B
HH -> N
CB -> H
NH -> C
HB -> C
HC -> B
HN -> C
NN -> C
BH -> H
NC -> B
NB -> B
BN -> B
BB -> N
BC -> B
CC -> N
CN -> C";

    #[test]
    fn part1_test() -> anyhow::Result<()> {
        let d = Day14::init(INPUT.to_string())?;
        assert_eq!("1588", d.part1()?);
        Ok(())
    }

    #[test]
    fn part2_test() -> anyhow::Result<()> {
        let d = Day14::init(INPUT.to_string())?;
        assert_eq!("2188189693529", d.part2()?);
        Ok(())
    }

    day_tests!(Day14, "14", "2590", "2875665202438");
}