use crate::prelude::*;

pub fn solution() -> Solution {
  Solution::new(Day::new(15, 2021), part_1, part_2).with_expected(386, 2806)
}

type Grid = HashMap<Coordinate, isize>;

#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
struct Coordinate(isize, isize);

impl Coordinate {
  fn distance(&self, target: &Coordinate) -> usize {
    self.0.abs_diff(target.0) + self.1.abs_diff(target.1)
  }

  fn successors(&self, grid: &Grid) -> Vec<(Coordinate, isize)> {
    let &Coordinate(x, y) = self;
    let mut successors = vec![];

    for coordinate in [
      Coordinate(x, y - 1), // Up
      Coordinate(x - 1, y), // Left
      Coordinate(x, y + 1), // Down
      Coordinate(x + 1, y), // Right
    ] {
      if let Some(value) = grid.get(&coordinate) {
        successors.push((coordinate, *value));
      }
    }

    successors
  }
}

fn parse(input: &str) -> Result<(Grid, Coordinate)> {
  let mut grid = Grid::new();
  let mut end = Coordinate(0, 0);

  for (y, line) in input.lines().enumerate() {
    let y = (y + 1).try_into()?;

    for (x, character) in line.char_indices() {
      let x = (x + 1).try_into()?;
      grid.insert(Coordinate(x, y), character.to_string().parse()?);

      if end.0 < x && end.1 < y {
        end = Coordinate(x, y);
      }
    }
  }

  Ok((grid, end))
}

fn enlarge_grid(grid: Grid, end: Coordinate) -> (Grid, Coordinate) {
  let Coordinate(width, height) = end;
  let mut larger_grid = grid.clone();

  for (Coordinate(x, y), risk) in grid {
    for step_x in 0..5 {
      for step_y in 0..5 {
        let risk = match (risk + step_x + step_y) % 9 {
          0 => 9,
          n => n,
        };
        let x = x + (width * step_x);
        let y = y + (height * step_y);
        larger_grid.insert(Coordinate(x, y), risk);
      }
    }
  }

  (larger_grid, Coordinate(width * 5, height * 5))
}

fn run(grid: Grid, end: Coordinate) -> Result<String> {
  Ok(
    astar(
      &Coordinate(1, 1),
      |p| p.successors(&grid),
      |p| p.distance(&end).try_into().unwrap(),
      |p| p == &end,
    )
    .ok_or_else(|| eyre!("No path found"))?
    .1
    .to_string(),
  )
}

fn part_1(input: &str) -> Result<String> {
  let (grid, end) = parse(input)?;
  run(grid, end)
}

fn part_2(input: &str) -> Result<String> {
  let (grid, end) = parse(input)?;
  let (grid, end) = enlarge_grid(grid, end);
  run(grid, end)
}