open Angstrom
open Base
open Core

type device = Empty | MirrorFwd | MirrorBwd | SplitterV | SplitterH
[@@deriving show]

type beam = int [@@deriving show]
type direction = North | East | South | West [@@deriving show]

let beam_has beam = function
  | North -> beam land 0b1000 <> 0
  | East -> beam land 0b0100 <> 0
  | South -> beam land 0b0010 <> 0
  | West -> beam land 0b0001 <> 0

let beam_set beam = function
  | North -> beam lor 0b1000
  | East -> beam lor 0b0100
  | South -> beam lor 0b0010
  | West -> beam lor 0b0001

let beam_energized beam = beam <> 0

type tile = { device : device; beam : beam } [@@deriving show]
type field = { tiles : tile array; rows : int; cols : int } [@@deriving show]

let rc_to_idx field r c = (r * field.cols) + c
let idx_to_rc field idx = (idx / field.cols, idx mod field.cols)
let field_get field r c = field.tiles.(rc_to_idx field r c)
let field_set field r c tile = field.tiles.(rc_to_idx field r c) <- tile

let field_mark field r c dir =
  let tile = field_get field r c in
  let beam = beam_set tile.beam dir in
  field_set field r c { tile with beam }

let print_field f = Out_channel.print_endline (show_field f)

module Parse = struct
  let parse_line s =
    let open Angstrom in
    let device = function
      | '.' -> return Empty
      | '/' -> return MirrorFwd
      | '\\' -> return MirrorBwd
      | '|' -> return SplitterV
      | '-' -> return SplitterH
      | _ -> fail "invalid device"
    and chars = char '.' <|> char '/' <|> char '\\' <|> char '|' <|> char '-' in
    let line = many1 (chars >>= device) <* end_of_input in
    parse_string ~consume:All line s

  let parse_input s =
    let lines = String.split_lines s in
    let rows = List.length lines
    and cols = String.length (List.hd_exn lines)
    and parsed = List.map lines ~f:parse_line in
    let devices = List.concat_map parsed ~f:Result.ok_or_failwith in
    let tiles = List.map ~f:(fun device -> { device; beam = 0 }) devices in
    { tiles = Array.of_list tiles; rows; cols }
end

module Part1 = struct
  type pos = int * int * direction
  type state = pos list

  let initial_state = [ (0, 0, East) ]

  (* return the next tile in the given direction, if it exists *)
  let next_tile field (r, c, dir) =
    let r', c' =
      match dir with
      | North -> (r - 1, c)
      | East -> (r, c + 1)
      | South -> (r + 1, c)
      | West -> (r, c - 1)
    in
    match (r', c') with
    | r, c when r >= 0 && r < field.rows && c >= 0 && c < field.cols ->
        Some (r, c, dir)
    | _ -> None

  (* step the beam forward one tile, returning a list of new positions.
     The list can have one element (for empty, mirror, some splitters) or
     two elements (for splitters). *)
  let step field ((r, c, dir) : pos) =
    let tile = field_get field r c in
    let dirs =
      match tile.device with
      | Empty -> [ dir ]
      | MirrorFwd -> (
          match dir with
          | North -> [ East ]
          | East -> [ North ]
          | South -> [ West ]
          | West -> [ South ])
      | MirrorBwd -> (
          match dir with
          | North -> [ West ]
          | East -> [ South ]
          | South -> [ East ]
          | West -> [ North ])
      | SplitterV -> (
          match dir with
          | East -> [ North; South ]
          | West -> [ North; South ]
          | d -> [ d ])
      | SplitterH -> (
          match dir with
          | North -> [ East; West ]
          | South -> [ East; West ]
          | d -> [ d ])
    in
    let new_beams = List.map dirs ~f:(fun d -> (r, c, d)) in
    let next_tiles = List.filter_map new_beams ~f:(next_tile field) in
    next_tiles

  (* traverse the field, marking tiles as visited. A list of positions is kept as state.
     A tile is visited at most four times, once for each direction. *)
  let rec traverse field : state -> unit = function
    | [] -> ()
    | ((r, c, dir) as pos) :: rest ->
        let tile = field_get field r c in
        let visited = beam_has tile.beam dir in
        if not visited then (
          field_mark field r c dir;
          traverse field (step field pos @ rest))
        else traverse field rest

  (* count the number of energized tiles in the field *)
  let count_energized field =
    let f sum tile = if beam_energized tile.beam then sum + 1 else sum in
    Array.fold field.tiles ~init:0 ~f
end

module Part2 = struct
  (* return a list of all edge tiles with according directions *)
  let all_edge_tiles field =
    let top = List.init field.cols ~f:(fun c -> (0, c, South))
    and bottom = List.init field.cols ~f:(fun c -> (field.rows - 1, c, North))
    and left = List.init field.rows ~f:(fun r -> (r, 0, East))
    and right = List.init field.rows ~f:(fun r -> (r, field.cols - 1, West)) in
    List.concat [ top; bottom; left; right ]

  (* reset the beam on all tiles *)
  let reset_field field =
    let f i tile = field.tiles.(i) <- { tile with beam = 0 } in
    Array.iteri field.tiles ~f

  (* count the number of energized tiles starting from the given position *)
  let count_energized field (init : Part1.pos) =
    reset_field field;
    Part1.traverse field [ init ];
    Part1.count_energized field

  (* find the maximum number of energized tiles starting from any edge tile *)
  let max_energized field =
    let f max init = Int.max max (count_energized field init) in
    List.fold (all_edge_tiles field) ~init:0 ~f
end

let () =
  let input = In_channel.input_all In_channel.stdin in
  let field = Parse.parse_input input in
  Part1.traverse field Part1.initial_state;
  let count = Part1.count_energized field in
  Out_channel.printf "Part 1: %d\n" count;
  let max = Part2.max_energized field in
  Out_channel.printf "Part 2: %d\n" max