let find grid c =
  match Aoc.Grid.idx_from_opt grid 0 c with
  | None -> failwith "find"
  | Some idx -> Aoc.Grid.pos_of_idx grid idx

let input_of_file fname =
  let grid = Aoc.Grid.of_file fname in
  let start_pos = find grid 'S' in
  (grid, start_pos)

let rec dijkstra visit check_end states =
  let compare_costs (lhs, _) (rhs, _) = compare lhs rhs in
  match states with
  | [] -> None
  | (cost, state) :: t ->
      if check_end state then Some (cost, state, t)
      else
        let new_states = visit cost state |> List.merge compare_costs t in
        dijkstra visit check_end new_states

let visited_idx grid ((dx, dy), p) =
  let add =
    match (dx, dy) with
    | 1, 0 -> 0
    | 0, 1 -> 1
    | -1, 0 -> 2
    | 0, -1 -> 3
    | _ -> failwith "visited_idx"
  in
  (Aoc.Grid.idx_of_pos grid p * 4) + add

let visit grid visited_grid cost state =
  let (dx, dy), ((x, y) as p) = List.hd state in
  let has_visited = visited_grid.(visited_idx grid (List.hd state)) in
  if has_visited then []
  else if Aoc.Grid.get_by_pos grid p = '#' then []
  else (
    visited_grid.(visited_idx grid (List.hd state)) <- true;
    [
      (cost + 1, ((dx, dy), (x + dx, y + dy)) :: state);
      (cost + 1000, ((-dy, dx), p) :: state);
      (cost + 1000, ((dy, -dx), p) :: state);
    ])

let has_visited grid visited_grid (cost, state) =
  visited_grid.(visited_idx grid (List.hd state)) < cost

let visit_max grid visited_grid cost state =
  let (dx, dy), ((x, y) as p) = List.hd state in
  if has_visited grid visited_grid (cost, state) then []
  else if Aoc.Grid.get_by_pos grid p = '#' then []
  else (
    Printf.printf "%d (%d, %d)\n" cost x y;
    flush stdout;
    visited_grid.(visited_idx grid (List.hd state)) <- cost;
    [
      (cost + 1, ((dx, dy), (x + dx, y + dy)) :: state);
      (cost + 1000, ((-dy, dx), p) :: state);
      (cost + 1000, ((dy, -dx), p) :: state);
    ]
    |> List.filter (fun x -> not (has_visited grid visited_grid x)))

let[@warning "-32"] print_point (x, y) = Printf.printf "(%d, %d)" x y

let check_end grid state =
  let _, p = List.hd state in
  Aoc.Grid.get_by_pos grid p = 'E'

let check_end2 grid _ state = check_end grid state

let part1 (grid, start_pos) =
  let visited_grid = Array.make (Aoc.Grid.length grid * 4) false in
  match
    dijkstra (visit grid visited_grid) (check_end grid)
      [ (0, [ ((1, 0), start_pos) ]) ]
  with
  | None -> failwith "part"
  | Some (cost, _, _) -> cost

let part2 (grid, start_pos) =
  let cost = part1 (grid, start_pos) in
  let visited_grid = Array.make (Aoc.Grid.length grid * 4) cost in
  let rec impl acc lst =
    match
      dijkstra (visit_max grid visited_grid) (check_end2 grid visited_grid) lst
    with
    | None -> acc
    | Some (_, states, remainder) ->
        List.iter (fun x -> visited_grid.(visited_idx grid x) <- 0) states;
        impl (states :: acc) remainder
  in
  impl [] [ (0, [ ((1, 0), start_pos) ]) ]
  |> List.concat |> List.map snd |> List.sort_uniq compare |> List.length

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