--- a/2023/day17.ml	Sat Dec 30 17:49:10 2023 +0100
+++ b/2023/day17.ml	Sun Dec 31 09:21:55 2023 +0100
@@ -37,12 +37,23 @@
     Sexp.List (List.map elts ~f:Elt.sexp_of_t)
 end
 
+module Int_pq = PrioQueue (Int)
+
+let () =
+  let q = Int_pq.empty in
+  let q = Int_pq.add q ~prio:1 1 in
+  let q = Int_pq.add q ~prio:2 2 in
+  let q = Int_pq.add q ~prio:0 3 in
+  assert (Int_pq.min_elt_exn q = 3)
+
 module type Comparable = sig
   type t [@@deriving sexp, compare]
 end
 
 type direction = North | East | South | West [@@deriving sexp, compare, eq]
-type heatloss = { value : int; min_so_far : int } [@@deriving sexp, compare]
+
+type heatloss = { value : int; min_so_far : int; prev : int * int }
+[@@deriving sexp, compare]
 
 type field = { r : int; c : int; field : heatloss array }
 [@@deriving sexp, compare]
@@ -58,7 +69,7 @@
 module Parse = struct
   let parse_tile c =
     String.of_char c |> Int.of_string |> fun value ->
-    { value; min_so_far = Int.max_value }
+    { value; min_so_far = Int.max_value; prev = (0, 0) }
 
   let parse_line s = String.to_list s |> List.map ~f:parse_tile
 
@@ -113,12 +124,14 @@
   let initial = [ Position.initial ]
   let max_straight = 3
 
+  (* Return potential neighbors at position (r, c) in direction dir *)
   let neighbors r c : direction -> neighbor list = function
     | North -> [ (r - 1, c, North); (r, c - 1, West); (r, c + 1, East) ]
     | East -> [ (r - 1, c, North); (r + 1, c, South); (r, c + 1, East) ]
     | South -> [ (r + 1, c, South); (r, c - 1, West); (r, c + 1, East) ]
     | West -> [ (r - 1, c, North); (r + 1, c, South); (r, c - 1, West) ]
 
+  (* Filter out neighbors that are not valid *)
   let valid_neighbors field Position.{ dir; straight; heatloss; _ } neighbors :
       neighbor list =
     let straight_ok = function
@@ -137,6 +150,7 @@
     in
     List.filter neighbors ~f:valid
 
+  (* From position pos, return a list of next tiles to go *)
   let next_options field (Position.{ r; c; dir; straight; _ } as pos) =
     let neighbors = neighbors r c dir in
     let neighbors' = valid_neighbors field pos neighbors in
@@ -147,12 +161,46 @@
           c = c';
           dir = dir';
           prev = (r, c);
-          straight = (if Int.(straight = 1) || equal_direction dir dir' then straight + 1 else 1);
+          straight =
+            (if Int.(straight = 1) || equal_direction dir dir' then straight + 1
+             else 1);
           heatloss = pos.heatloss + (field_at field r' c').value;
         }
     in
     List.map neighbors' ~f:pos_of_neighbor
 
+  let rec trace_path ?(acc = []) field r c =
+    match (r, c) with
+    | 0, 0 -> (0, 0) :: acc
+    | _ ->
+        let entry = field_at field r c in
+        let r', c' = entry.prev in
+        trace_path ~acc:((r, c) :: acc) field r' c'
+
+  let start_trace_path field =
+    let r, c = dst field in
+    trace_path field r c
+
+  let visualizer_field field path =
+    let a = Array.map field.field ~f:(fun _ -> 0) in
+    let f (r, c) = a.(rc_to_ix field r c) <- 1 in
+    List.iter path ~f;
+    a
+
+  let print_visualizer_field field r c =
+    let rows = Sequence.range 0 r and cols = Sequence.range 0 c in
+    let f r' =
+      Sequence.iter cols ~f:(fun c' -> printf "%d" field.((r' * r) + c'));
+      printf "\n"
+    in
+    Sequence.iter rows ~f
+
+  let visualize_field field path =
+    let f = visualizer_field field path in
+    print_visualizer_field f field.r field.c
+
+  (* Apply Dijkstra's algorithm (or something like that...)
+     to find the shortest path according to restrictions. *)
   let solve field =
     let dstr, dstc = dst field in
     let rec loop (q : Pospq.t) =
@@ -160,18 +208,14 @@
       else
         let q, pos = Pospq.remove_min_elt_exn q in
         if Int.equal pos.r dstr && Int.equal pos.c dstc then (
-          Out_channel.printf "Visiting %s\n"
-            (Position.sexp_of_t pos |> Sexp.to_string_hum);
+          field_update field pos.r pos.c (fun v ->
+              { v with min_so_far = pos.heatloss; prev = pos.prev });
           Some pos.heatloss)
         else
           let min_so_far = (field_at field pos.r pos.c).min_so_far in
           if min_so_far = Int.max_value then (
-            Out_channel.printf "Visiting %s (%d vs. %d)\n"
-              (Position.sexp_of_t pos |> Sexp.to_string_hum)
-              pos.heatloss min_so_far;
-
             field_update field pos.r pos.c (fun v ->
-                { v with min_so_far = pos.heatloss });
+                { v with min_so_far = pos.heatloss; prev = pos.prev });
             let next = next_options field pos in
             let q =
               List.fold_left next ~init:q ~f:(fun q' opt ->
@@ -187,4 +231,11 @@
   let inp = In_channel.(input_all stdin) in
   let field = Parse.parse_field inp in
   let part1 = Option.value_exn (Part1.solve field) in
-  Out_channel.printf "Part1: %d\n" part1
+  let trace = Part1.start_trace_path field in
+  Out_channel.printf "Part1: %d\n" part1;
+  Out_channel.printf "Path: %s\n"
+    (Sexp.to_string_hum
+       (List.sexp_of_t
+          (fun (r, c) -> Sexp.List [ Int.sexp_of_t r; Int.sexp_of_t c ])
+        trace));
+  Part1.visualize_field field trace