type map = { map : string; width : int; height : int }
(** [map] represents a map with layout [map] and a [width] and [height].
    [map.map] is a string of length [width * height] with the [i]th character
    describing the position [(i mod map.width, i / map.width)]. *)

(** [find_start strs] returns the location [(x, y)] of the starting position. *)
let find_start map =
  match String.index_from_opt map.map 0 '^' with
  | Some i -> (i mod map.width, i / map.width)
  | None -> failwith "find_start"

(** [map_of_strings strs] returns the map generated from the input map described
    by the list of strings [strs] *)
let map_of_strings strs =
  let width = String.length (List.hd strs) in
  let map = List.fold_left String.cat "" strs in
  let height = String.length map / width in
  { map; width; height }

(** [read_file fname] reads the input map from [fname]. It returns a
    [(map, pos, vel)] tuple, consisting of the obsticle map, initial position,
    and initial velocity. *)
let read_file fname =
  let lst = Aoc.strings_of_file fname in
  let map = map_of_strings lst in
  let pos = find_start map in
  (map, pos, (0, -1))

(** [is_valid_pos map pos] returns true if the position [pos] is valid for the
    map [map]. *)
let is_valid_pos map (x, y) =
  if y < 0 || y >= map.height then false
  else if x < 0 || x >= map.width then false
  else true

(** [is_block map pos] returns [true] iff the location [pos] is a blockages in
    [map]. *)
let is_block map (x, y) =
  if is_valid_pos map (x, y) then map.map.[x + (y * map.width)] = '#' else false

(** [insert_block map pos] inserts a blockage at [pos] into [map] and returns
    the new map. *)
let insert_block map (x, y) =
  let idx = x + (y * map.width) in
  let start = String.sub map.map 0 idx in
  let e = String.sub map.map (idx + 1) (String.length map.map - idx - 1) in
  { map with map = start ^ "#" ^ e }

(** [move map (pos, vel)] moves [pos] one step forward on the [map]. [vel] gives
    the movement vector. If the movement will cause an obstacle to be hit then
    [vel] is rotated right by 90 degrees and we move in that direction. Returns
    the updated [(pos, vel)] pair. *)
let rec move map ((x, y), (dx, dy)) =
  let x', y' = (x + dx, y + dy) in
  if is_block map (x', y') then move map ((x, y), (-dy, dx))
  else ((x', y'), (dx, dy))

(** [walk_map map (pos, vel)] walks around [map] starting at [pos] moving in the
    direction [vel]. It returns a list of all positions visited before falling
    off one of the sides. *)
let walk_map map (pos, vel) =
  let rec impl acc (pos, vel) =
    if is_valid_pos map pos then impl (pos :: acc) (move map (pos, vel))
    else acc
  in
  impl [] (pos, vel)

(** [has_cycles map (pos, vel)] returns true if walking around [map] starting at
    [pos] going in [vel] direction will end up in a never ending cycle.*)
let has_cycles map start =
  (* We detect a cycle by walking two 'agents' around the map from the same
     starting position.  Agent 1 moves 1 step at a time, agent 2 moves 2.  If
     the agents ever end up on the same square facing the same direction we have
     a cycle.  This works even if the cycle doesn't start immediately. *)
  let rec impl agent1 ((pos', _) as agent2) =
    (* Only need to check pos' for validity because if pos is not valid then
       pos' must also be invalid, and have been invalid before this. *)
    if not (is_valid_pos map pos') then false
    else if agent1 = agent2 then true
    else impl (move map agent1) (move map (move map agent2))
  in
  (* Start Agent 2 a step ahead of Agent 1 so we don't fail at the start 
     position. *)
  impl start (move map start)

(** [walk_block map (pos, vel) bpos] adds a block to the map [map] at [bpos] and
    then sees if walking the map starting with [(pos, vel)] has a cycle. *)
let walk_block map (pos, vel) bpos =
  if bpos = pos then false
  else
    let map' = insert_block map bpos in
    has_cycles map' (pos, vel)

let part1 (map, pos, vel) =
  walk_map map (pos, vel) |> List.sort_uniq Aoc.IntPair.compare |> List.length

let part2 (map, pos, vel) =
  walk_map map (pos, vel)
  |> List.sort_uniq Aoc.IntPair.compare
  |> List.filter (walk_block map (pos, vel))
  |> List.length

let _ = Aoc.main read_file [ (string_of_int, part1); (string_of_int, part2) ]