AoC-2024/src/days/day6.rs

use std::path::Path;

use advent_derive::advent_day;
use itertools::Itertools;

use crate::{error::AdventError, input::get_lines};
use super::*;

#[derive(PartialEq, Clone, Copy, PartialOrd, Eq)]
enum Direction {
    Up,
    Right,
    Down,
    Left
}

impl Ord for Direction {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        match self {
            Direction::Up => {
                match other {
                    Direction::Up => std::cmp::Ordering::Equal,
                    Direction::Right => std::cmp::Ordering::Greater,
                    Direction::Down => std::cmp::Ordering::Greater,
                    Direction::Left => std::cmp::Ordering::Greater,
                }
            },
            Direction::Right => {
                match other {
                    Direction::Up => std::cmp::Ordering::Less,
                    Direction::Right => std::cmp::Ordering::Equal,
                    Direction::Down => std::cmp::Ordering::Greater,
                    Direction::Left => std::cmp::Ordering::Greater,
                }
            },
            Direction::Down => {
                match other {
                    Direction::Up => std::cmp::Ordering::Less,
                    Direction::Right => std::cmp::Ordering::Less,
                    Direction::Down => std::cmp::Ordering::Equal,
                    Direction::Left => std::cmp::Ordering::Greater,
                }
            },
            Direction::Left => {
                match other {
                    Direction::Up => std::cmp::Ordering::Less,
                    Direction::Right => std::cmp::Ordering::Less,
                    Direction::Down => std::cmp::Ordering::Less,
                    Direction::Left => std::cmp::Ordering::Equal,
                }
            },
        }
    }
}

struct Guard {
    position: (i32, i32),
    direction: Direction
}

struct Day6Part1 {
    obstacles: Vec<(i32, i32)>,
    guard: Guard,
    width: usize,
    height: usize
}

impl Day6Part1 {
    fn new() -> Self {
        Self {
            obstacles: vec![],
            guard: Guard {
                position: (0, 0),
                direction: Direction::Up
            },
            width: 0,
            height: 0
        }
    }
}

struct Day6Part2 {
    obstacles: Vec<(i32, i32)>,
    guard: Guard,
    width: usize,
    height: usize
}

impl Day6Part2 {
    fn new() -> Self {
        Self {
            obstacles: vec![],
            guard: Guard {
                position: (0, 0),
                direction: Direction::Up
            },
            width: 0,
            height: 0
        }
    }
}

fn read_input_part1(part1: &mut Day6Part1, path: &Path) -> Result<(), AdventError> {    
    let iter = get_lines(path).enumerate();
    for (y, s) in iter {
        let s = s?;
        part1.width = s.len();
        part1.height = y+1;
        for (x, _) in s.chars().enumerate().filter(|(_, c)| { *c == '#'}) {
            part1.obstacles.push((x as i32, y as i32));
        }
        if let Some(x) = s.chars().position(|c| { c == '^'}) {
            part1.guard.position = (x as i32, y as i32);
        }
    }
    Ok(())
}

fn solve_part1(part1: &mut Day6Part1) -> Result<u64, AdventError> {
    let mut spaces = vec![];
    spaces.push(part1.guard.position);
    loop {
        if let Some(x) = step(part1) {
            spaces.push(x);
        } else {
            break;
        }
    }
    Ok(spaces.iter().unique().count() as u64)
}

fn step(part1: &mut Day6Part1) -> Option<(i32, i32)> {
    let next_space = match part1.guard.direction {
        Direction::Up => (part1.guard.position.0, part1.guard.position.1 - 1),
        Direction::Right => (part1.guard.position.0 + 1, part1.guard.position.1),
        Direction::Down => (part1.guard.position.0, part1.guard.position.1 + 1),
        Direction::Left => (part1.guard.position.0 - 1, part1.guard.position.1)
    };

    if part1.obstacles.contains(&next_space) {
        part1.guard.direction = match part1.guard.direction {
            Direction::Up => Direction::Right,
            Direction::Right => Direction::Down,
            Direction::Down => Direction::Left,
            Direction::Left => Direction::Up,
        };
        return step(part1);
    }

    if next_space.0 < 0 || next_space.0 > part1.width as i32 || next_space.1 < 0 || next_space.1 >= part1.height as i32 {
        None
    } else {
        part1.guard.position = next_space;
        Some(next_space)
    }
}

fn read_input_part2(part2: &mut Day6Part2, path: &Path) -> Result<(), AdventError> {
    let iter = get_lines(path).enumerate();
    for (y, s) in iter {
        let s = s?;
        part2.width = s.len();
        part2.height = y+1;
        for (x, _) in s.chars().enumerate().filter(|(_, c)| { *c == '#'}) {
            part2.obstacles.push((x as i32, y as i32));
        }
        if let Some(x) = s.chars().position(|c| { c == '^'}) {
            part2.guard.position = (x as i32, y as i32);
        }
    }
    Ok(())
}

fn step2(part2: &mut Day6Part2) -> Option<(i32, i32)> {
    let next_space = match part2.guard.direction {
        Direction::Up => (part2.guard.position.0, part2.guard.position.1 - 1),
        Direction::Right => (part2.guard.position.0 + 1, part2.guard.position.1),
        Direction::Down => (part2.guard.position.0, part2.guard.position.1 + 1),
        Direction::Left => (part2.guard.position.0 - 1, part2.guard.position.1)
    };

    if let Ok(_) = part2.obstacles.binary_search(&next_space) {
        part2.guard.direction = match part2.guard.direction {
            Direction::Up => Direction::Right,
            Direction::Right => Direction::Down,
            Direction::Down => Direction::Left,
            Direction::Left => Direction::Up,
        };
        return step2(part2);
    }

    if next_space.0 < 0 || next_space.0 > part2.width as i32 || next_space.1 < 0 || next_space.1 >= part2.height as i32 {
        None
    } else {
        part2.guard.position = next_space;
        Some(next_space)
    }
}

fn solve_part2(part2: &mut Day6Part2) -> Result<u64, AdventError> {
    let start_pos = part2.guard.position;
    let start_dir = part2.guard.direction;

    part2.obstacles.sort();
    let mut spaces = vec![];
    loop {
        if let Some(x) = step2(part2) {
            spaces.push(x);
        } else {
            break;
        }
    }

    let obstacles_to_be = spaces.into_iter().unique().collect::<Vec<(i32, i32)>>();
    let mut loops = 0u64;
    for item in obstacles_to_be.into_iter() {
        let mut been = vec![];
        part2.guard.position = start_pos;
        part2.guard.direction = start_dir;
        let remove_index;
        match part2.obstacles.binary_search(&item) {
            Ok(_) => {
                return Err(AdventError("Should not happen".to_string()));
            }
            Err(pos) => {
                remove_index = pos;
                part2.obstacles.insert(pos, item)
            }
        }

        loop {
            if let Some(pos) = step2(part2) {
                let needle = (pos, part2.guard.direction);
                if let Ok(_) = been.binary_search(&(pos, part2.guard.direction)) {
                    loops += 1;
                    break;
                }
                match been.binary_search(&needle) {
                    Ok(_) => {
                        return Err(AdventError("Should not happen".to_string()));
                    }
                    Err(pos) => {
                        been.insert(pos, needle)
                    }
                }
            } else {
                break;
            }
        }
        part2.obstacles.remove(remove_index);
    }

    Ok(loops)
    
}

#[advent_day(Day6Part1, Day6Part2)]
struct Day6;