Day 16: The Floor Will Be Lava

use std::fs;
use std::path::PathBuf;

use clap::Parser;

use rayon::prelude::*;

#[derive(Parser)]
#[command(author, version, about, long_about = None)]
struct Cli {
    input_file: PathBuf,
}

#[derive(Copy, Clone)]
enum TileState {
    None,
    Energized(BeamState),
}
#[derive(Default, Copy, Clone)]
struct BeamState {
    up: bool,
    down: bool,
    left: bool,
    right: bool,
}

fn main() {
    // Parse CLI arguments
    let cli = Cli::parse();

    // Read file
    let input_text = fs::read_to_string(&cli.input_file)
        .expect(format!("File \"{}\" not found", cli.input_file.display()).as_str());

    let tiles: Vec> = input_text.lines().map(|l| l.chars().collect()).collect();

    // Part 1
    let part_1 = test_beam(&tiles, (0, 0), (0, 1));
    println!("Part 1: {}", part_1);

    // Part 2
    let part_2: usize = (0..4)
        .into_par_iter()
        .map(|dir| {
            (0..tiles.len())
                .into_par_iter()
                .map(move |x| (dir.clone(), x))
        })
        .flatten()
        .map(|(dir, x)| match dir {
            0 => ((0, x), (1, 0)),
            1 => ((x, tiles[0].len() - 1), (0, -1)),
            2 => ((tiles.len() - 1, x), (-1, 0)),
            3 => ((x, 0), (0, 1)),
            _ => unreachable!(),
        })
        .map(|(loc, dir)| test_beam(&tiles, loc, dir))
        .max()
        .unwrap();
    println!("Part 2: {}", part_2);
}

fn test_beam(
    tiles: &Vec>,
    start_location: (usize, usize),
    start_direction: (i64, i64),
) -> usize {
    let mut energized: Vec> =
        vec![vec![TileState::None; tiles[0].len()]; tiles.len()];

    continue_beam(
        &mut energized,
        &tiles,
        start_location,
        start_direction,
        true,
        0,
    );
    energized
        .iter()
        .map(|r| {
            r.iter()
                .filter(|t| matches!(t, TileState::Energized(_)))
                .count()
        })
        .sum()
}

fn continue_beam(
    energized: &mut Vec>,
    tiles: &Vec>,
    beam_location: (usize, usize),
    beam_direction: (i64, i64),
    start_hack: bool,
    depth: usize,
) {
    assert_ne!(beam_direction, (0, 0));

    // Set current tile to energized with the direction
    let current_state = energized[beam_location.0][beam_location.1];
    if !start_hack {
        energized[beam_location.0][beam_location.1] = match current_state {
            TileState::None => TileState::Energized(match beam_direction {
                (0, 1) => BeamState {
                    right: true,
                    ..BeamState::default()
                },
                (0, -1) => BeamState {
                    left: true,
                    ..BeamState::default()
                },
                (1, 0) => BeamState {
                    down: true,
                    ..BeamState::default()
                },
                (-1, 0) => BeamState {
                    up: true,
                    ..BeamState::default()
                },
                _ => unreachable!(),
            }),
            TileState::Energized(state) => TileState::Energized(match beam_direction {
                (0, 1) => {
                    if state.right {
                        return;
                    }
                    BeamState {
                        right: true,
                        ..state
                    }
                }
                (0, -1) => {
                    if state.left {
                        return;
                    }
                    BeamState {
                        left: true,
                        ..state
                    }
                }
                (1, 0) => {
                    if state.down {
                        return;
                    }
                    BeamState {
                        down: true,
                        ..state
                    }
                }
                (-1, 0) => {
                    if state.up {
                        return;
                    }
                    BeamState { up: true, ..state }
                }
                _ => unreachable!(),
            }),
        };
    }

    // energized[beam_location.0][beam_location.1] = TileState::Energized(BeamState { up: , down: , left: , right:  });

    let next_beam_location = {
        let loc = (
            (beam_location.0 as i64 + beam_direction.0),
            (beam_location.1 as i64 + beam_direction.1),
        );

        if start_hack {
            beam_location
        } else if loc.0 < 0
            || loc.0 >= tiles.len() as i64
            || loc.1 < 0
            || loc.1 >= tiles[0].len() as i64
        {
            return;
        } else {
            (loc.0 as usize, loc.1 as usize)
        }
    };
    let next_beam_tile = tiles[next_beam_location.0][next_beam_location.1];

    let next_beam_directions: Vec<(i64, i64)> = match next_beam_tile {
        '.' => vec![beam_direction],
        '/' => match beam_direction {
            (0, 1) => vec![(-1, 0)],
            (0, -1) => vec![(1, 0)],
            (1, 0) => vec![(0, -1)],
            (-1, 0) => vec![(0, 1)],
            _ => unreachable!(),
        },
        '\\' => match beam_direction {
            (0, 1) => vec![(1, 0)],
            (0, -1) => vec![(-1, 0)],
            (1, 0) => vec![(0, 1)],
            (-1, 0) => vec![(0, -1)],
            _ => unreachable!(),
        },
        '|' => match beam_direction {
            (0, 1) => vec![(1, 0), (-1, 0)],
            (0, -1) => vec![(1, 0), (-1, 0)],
            (1, 0) => vec![(1, 0)],
            (-1, 0) => vec![(-1, 0)],
            _ => unreachable!(),
        },
        '-' => match beam_direction {
            (0, 1) => vec![(0, 1)],
            (0, -1) => vec![(0, -1)],
            (1, 0) => vec![(0, 1), (0, -1)],
            (-1, 0) => vec![(0, 1), (0, -1)],
            _ => unreachable!(),
        },
        _ => unreachable!(),
    };

    for dir in next_beam_directions {
        continue_beam(energized, tiles, next_beam_location, dir, false, depth + 1);
    }
}

26.28 line-seconds