use crate::prelude::*;

type FishMap = HashMap<isize, isize>;

pub fn solution() -> Solution {
  Solution::new(Day::new(6, 2021), part_1, part_2)
    .with_expected(343441, 1569108373832_i64)
}

fn part_1(input: &str) -> Result<String> {
  let mut fishes = parse_fishes(input)?;

  for _ in 0..80 {
    fishes = simulate_fishes(fishes);
  }

  Ok(count_fishes(fishes).to_string())
}

fn part_2(input: &str) -> Result<String> {
  let mut fishes = parse_fishes(input)?;

  for _ in 0..256 {
    fishes = simulate_fishes(fishes);
  }

  Ok(count_fishes(fishes).to_string())
}

fn parse_fishes(input: &str) -> Result<FishMap> {
  let mut fishes = FishMap::new();
  let individual_fishes = input
    .split(',')
    .map(str::parse)
    .collect::<Result<Vec<_>, _>>()?;

  for fish in individual_fishes {
    let amount = fishes.entry(fish).or_default();
    *amount += 1;
  }

  Ok(fishes)
}

fn count_fishes(fishes: FishMap) -> isize {
  fishes.into_iter().map(|(_, amount)| amount).sum()
}

fn simulate_fishes(fishes: FishMap) -> FishMap {
  let mut new_fishes = HashMap::new();

  for (mut timer, amount) in fishes {
    timer -= 1;

    if timer < 0 {
      let new_fish = new_fishes.entry(8).or_default();
      *new_fish += amount;

      let reset_fish = new_fishes.entry(6).or_default();
      *reset_fish += amount;
    } else {
      let current_fish = new_fishes.entry(timer).or_default();
      *current_fish += amount;
    }
  }

  new_fishes
}