import Control.Monad
import Data.Array.Unboxed (UArray)
import qualified Data.Array.Unboxed as A
import Data.Foldable
import Data.Set (Set)
import qualified Data.Set as Set

type Pos = (Int, Int)

readInput :: String -> UArray Pos Char
readInput s =
  let rows = lines s
   in A.listArray ((1, 1), (length rows, length $ head rows)) $ concat rows

energized :: (Pos, Pos) -> UArray Pos Char -> Set Pos
energized start grid = go Set.empty $ Set.singleton start
  where
    go seen beams
      | Set.null beams = Set.map fst seen
      | otherwise =
          let seen' = seen `Set.union` beams
              beams' = Set.fromList $ do
                ((y, x), (dy, dx)) <- toList beams
                d'@(dy', dx') <- case grid A.! (y, x) of
                  '/' -> [(-dx, -dy)]
                  '\\' -> [(dx, dy)]
                  '|' | dx /= 0 -> [(-1, 0), (1, 0)]
                  '-' | dy /= 0 -> [(0, -1), (0, 1)]
                  _ -> [(dy, dx)]
                let p' = (y + dy', x + dx')
                    beam' = (p', d')
                guard $ A.inRange (A.bounds grid) p'
                guard $ beam' `Set.notMember` seen'
                return beam'
           in go seen' beams'

part1 = Set.size . energized ((1, 1), (0, 1))

part2 input = maximum counts
  where
    (_, (h, w)) = A.bounds input
    starts =
      concat $
        [[((y, 1), (0, 1)), ((y, w), (0, -1))] | y <- [1 .. h]]
          ++ [[((1, x), (1, 0)), ((h, x), (-1, 0))] | x <- [1 .. w]]
    counts = map (\s -> Set.size $ energized s input) starts

main = do
  input <- readInput <$> readFile "input16"
  print $ part1 input
  print $ part2 input

A whopping 130.050 line-seconds!

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_iter()
        .map(|dir| (0..tiles.len()).into_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);
    }
}