open Angstrom
open Base
open Core

module PrioQueue (Inner : Comparator.S) = struct
  (* Element added to set *)
  module Elt = struct
    type t = int * Inner.t

    let compare (a, x) (b, y) =
      match Int.compare a b with 0 -> Inner.comparator.compare x y | x -> x

    let sexp_of_t (a, x) =
      Sexp.List [ Int.sexp_of_t a; Inner.comparator.sexp_of_t x ]
  end

  (* Comparator for Set *)
  module EltComp = struct
    include Comparator.Make (Elt)

    type t = Elt.t
  end

  type t = (Elt.t, EltComp.comparator_witness) Set.t

  let empty : t = Set.empty (module EltComp)
  let is_empty = Set.is_empty
  let add q ~prio elt = Set.add q (prio, elt)
  let min_elt_exn q = snd (Set.min_elt_exn q)

  let remove_min_elt_exn q =
    let ((_, elt) as min) = Set.min_elt_exn q in
    (Set.remove q min, elt)

  let sexp_of_t q =
    let elts = Set.to_list q in
    Sexp.List (List.map elts ~f:Elt.sexp_of_t)
end

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 field = { r : int; c : int; field : heatloss array }
[@@deriving sexp, compare]

let rc_to_ix field r c = (r * field.c) + c
let ix_to_rc field ix = Int.(ix / field.c, ix % field.c)
let field_at field r c = field.field.(rc_to_ix field r c)

let field_update field r c f =
  let v = field_at field r c in
  field.field.(rc_to_ix field r c) <- f v

module Parse = struct
  let parse_tile c =
    String.of_char c |> Int.of_string |> fun value ->
    { value; min_so_far = Int.max_value }

  let parse_line s = String.to_list s |> List.map ~f:parse_tile

  let parse_field inp =
    let lines = String.split inp ~on:'\n' in
    let r = List.length lines in
    let c = String.length (List.hd_exn lines) in
    let last = List.last_exn lines in
    let r = if String.is_empty last then r - 1 else r in
    let parsed_lines = List.map lines ~f:parse_line in
    let field = Array.of_list (List.concat parsed_lines) in
    { r; c; field }
end

module Position = struct
  type t = {
    r : int;
    c : int;
    prev : int * int;
    dir : direction;
    straight : int;
    heatloss : int;
  }
  [@@deriving sexp]

  let compare { heatloss = hl1; r = r1; c = c1; _ }
      { heatloss = hl2; r = r2; c = c2; _ } =
    match Int.compare hl1 hl2 with
    | 0 -> ( match Int.compare r1 r2 with 0 -> Int.compare c1 c2 | c -> c)
    | c -> c

  let initial =
    { r = 0; c = 0; dir = East; prev = (0, 0); straight = 0; heatloss = 0 }
end

(* Directly create a priority queue from a module containing compare/sexp_of_t *)
module PrimPrioQueue (Inner : Comparable) = struct
  module CompPos = struct
    include Inner
    include Comparable.Make (Inner)
  end

  include PrioQueue (CompPos)
end

module Pospq = PrimPrioQueue (Position)

module Part1 = struct
  type neighbor = int * int * direction

  let dst field = (field.r - 1, field.c - 1)
  let initial = [ Position.initial ]
  let max_straight = 3

  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) ]

  let valid_neighbors field Position.{ dir; straight; heatloss; _ } neighbors :
      neighbor list =
    let straight_ok = function
      | _, _, dir' when equal_direction dir dir' -> straight <= max_straight
      | _ -> true
    and within_field (r, c, _) = r >= 0 && r < field.r && c >= 0 && c < field.c
    and better_path (r, c, _) =
      let best_heatloss = (field_at field r c).min_so_far in
      let this_heatloss = heatloss + (field_at field r c).value in
      this_heatloss <= best_heatloss
    in
    let valid n =
      match straight_ok n && within_field n with
      | true -> better_path n
      | false -> false
    in
    List.filter neighbors ~f:valid

  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
    let pos_of_neighbor (r', c', dir') =
      Position.
        {
          r = r';
          c = c';
          dir = dir';
          prev = (r, c);
          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 solve field =
    let dstr, dstc = dst field in
    let rec loop (q : Pospq.t) =
      if Pospq.is_empty q then None
      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);
          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 });
            let next = next_options field pos in
            let q =
              List.fold_left next ~init:q ~f:(fun q' opt ->
                  Pospq.add q' ~prio:opt.heatloss opt)
            in
            loop q)
          else loop q
    in
    loop (List.fold_left initial ~init:Pospq.empty ~f:(Pospq.add ~prio:0))
end

let () =
  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