use nalgebra::{DMatrix, Point2, Unit, Vector2};
use nom::{
    character::complete::{line_ending, satisfy},
    combinator::{map, map_res},
    multi::{many1, separated_list1},
    IResult,
};
use std::{collections::HashSet, fs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_17.txt")?;
    let parsed = HeatlossMap::parser(&input).unwrap().1;
    dbg!(&parsed.find_shortest_heatloss_path(false));
    dbg!(&parsed.find_shortest_heatloss_path(true));

    Ok(())
}

#[derive(Debug)]
struct HeatlossMap(DMatrix<u32>);

#[derive(Debug, Hash, PartialEq, Eq, Clone)]
struct Position {
    pos: Point2<isize>,
    facing: Unit<Vector2<isize>>,
    duration_with_facing: usize,
}

impl HeatlossMap {
    fn parser(input: &str) -> IResult<&str, Self> {
        map(
            separated_list1(
                line_ending,
                many1(map_res(satisfy(|c| c.is_digit(10)), |d| {
                    d.to_string().parse::<u32>()
                })),
            ),
            |digits_vec| {
                HeatlossMap(DMatrix::from_iterator(
                    digits_vec.len(),
                    digits_vec[0].len(),
                    digits_vec.into_iter().flat_map(|row| row.into_iter()),
                ))
            },
        )(input)
    }

    fn find_shortest_heatloss_path(&self, is_ultra_crucible: bool) -> u32 {
        let start = Position {
            pos: Point2::new(0, 0),
            facing: Vector2::x_axis(),
            duration_with_facing: 0,
        };
        let end_pos = self.end();

        let mut frontier = vec![(0, start.clone())];
        let mut visited = HashSet::new();
        visited.insert(start);

        let mut distance_to_end = None;
        while distance_to_end.is_none() && frontier.len() > 0 {
            // shortest distance is now at the end
            frontier
                .sort_unstable_by(|(distance_a, _), (distance_b, _)| distance_b.cmp(distance_a));
            let (front_distance, front_position) = frontier.pop().unwrap();

            let mut next_points: Vec<Position> = Vec::new();

            if !is_ultra_crucible || front_position.duration_with_facing >= 4 {
                let facing_l = rotate_left(&front_position.facing);
                let pos_l = front_position.pos + *facing_l;
                next_points.push(Position {
                    pos: pos_l,
                    facing: facing_l,
                    duration_with_facing: 1,
                });

                let facing_r = rotate_right(&front_position.facing);
                let pos_r = front_position.pos + *facing_r;
                next_points.push(Position {
                    pos: pos_r,
                    facing: facing_r,
                    duration_with_facing: 1,
                });
            }

            if (is_ultra_crucible && front_position.duration_with_facing < 10)
                || (!is_ultra_crucible && front_position.duration_with_facing < 3)
            {
                let pos = front_position.pos + *front_position.facing;
                next_points.push(Position {
                    pos,
                    facing: front_position.facing.clone(),
                    duration_with_facing: front_position.duration_with_facing + 1,
                });
            }

            for next_point in next_points {
                if let Some(step_distance) = self.get(&next_point.pos) {
                    if !visited.contains(&next_point) {
                        visited.insert(next_point.clone());

                        let distance = front_distance + step_distance;

                        if next_point.pos == end_pos {
                            if is_ultra_crucible {
                                if next_point.duration_with_facing >= 4 {
                                    distance_to_end = Some(distance);
                                }
                            } else {
                                distance_to_end = Some(distance);
                            }
                        }
                        frontier.push((distance, next_point));
                    }
                }
            }
        }

        distance_to_end.unwrap()
    }

    fn end(&self) -> Point2<isize> {
        Point2::new(self.0.ncols() as isize - 1, self.0.nrows() as isize - 1)
    }

    fn get(&self, pos: &Point2<isize>) -> Option<u32> {
        let x: Option<usize> = pos.x.try_into().ok();
        let y: Option<usize> = pos.y.try_into().ok();

        match (x, y) {
            (Some(x), Some(y)) => self.0.get((x, y)).copied(),
            _ => None,
        }
    }
}

fn rotate_left(facing: &Unit<Vector2<isize>>) -> Unit<Vector2<isize>> {
    Unit::new_unchecked(Vector2::new(-facing.y, facing.x))
}

fn rotate_right(facing: &Unit<Vector2<isize>>) -> Unit<Vector2<isize>> {
    Unit::new_unchecked(Vector2::new(facing.y, -facing.x))
}