use crate::prelude::*;

type PairMap = HashMap<(char, char), char>;
type PairCounts = HashMap<(char, char), isize>;

pub fn solution() -> Solution {
  Solution::new(Day::new(14, 2021), part_1, part_2)
    .with_expected(3411, 7477815755570_i64)
}

fn parse(input: &str) -> Result<(String, PairMap)> {
  let mut lines = input.lines();
  let template = lines
    .next()
    .ok_or_else(|| eyre!("Invalid input: {}", input))?
    .to_string();

  let mut pairs = HashMap::new();

  for line in lines.skip(1) {
    let (a, b, c) = line
      .replace(" -> ", "")
      .chars()
      .collect_tuple()
      .ok_or_else(|| eyre!("Invalid line: {}", line))?;

    pairs.insert((a, b), c);
  }

  Ok((template, pairs))
}

fn apply(counts: &PairCounts, pairs: &PairMap) -> PairCounts {
  let mut to_add = PairCounts::new();

  for (pair, count) in counts {
    if let Some(character) = pairs.get(pair) {
      *to_add.entry((pair.0, *character)).or_default() += count;
      *to_add.entry((*character, pair.1)).or_default() += count;
    }
  }

  to_add
}

fn count_totals(counts: &PairCounts) -> HashMap<char, isize> {
  let mut totals = HashMap::new();

  for ((_, b), count) in counts {
    // If I count the first character in the pair the end result ends up being
    // off by one...?
    *totals.entry(*b).or_default() += count;
  }

  totals
}

fn run(input: &str, steps: isize) -> Result<String> {
  let (template, pairs) = parse(input)?;

  let mut counts = PairCounts::new();
  for (a, b) in template.chars().tuple_windows() {
    *counts.entry((a, b)).or_default() += 1;
  }

  for _ in 0..steps {
    counts = apply(&counts, &pairs);
  }

  let totals = count_totals(&counts)
    .into_iter()
    .sorted_by(|(_, a), (_, b)| a.cmp(b));

  let (_, min) = totals
    .clone()
    .next()
    .ok_or_else(|| eyre!("No minimum found"))?;
  let (_, max) = totals.last().ok_or_else(|| eyre!("No maximum found"))?;

  Ok((max - min).to_string())
}

fn part_1(input: &str) -> Result<String> {
  run(input, 10)
}

fn part_2(input: &str) -> Result<String> {
  run(input, 40)
}