open Base
open Core
open Angstrom

module Position = struct
  (* x, y starting at top left corner. x is row, y is column. *)
  type t = Position of int * int [@@deriving sexp, compare]

  let hash (Position (x, y)) =
    let h = Hash.create () in
    let h = Hash.fold_int h x in
    let h = Hash.fold_int h y in
    Hash.get_hash_value h

  let eq a b = 0 = compare a b
end

type pipe = Vert | Horiz | NE | NW | SW | SE | Ground | Start
[@@deriving sexp, compare]

type dir = N | E | S | W [@@deriving sexp]
type field = { rows : int; cols : int; pipes : (Position.t, pipe) Hashtbl.t }

let find_start field =
  let f ~key ~data acc =
    match acc with
    | None -> if 0 = compare_pipe data Start then Some key else None
    | x -> x
  in
  let maybe_start = Hashtbl.fold ~init:None ~f field in
  Option.value_exn maybe_start

let make_field rows cols =
  { pipes = Hashtbl.create (module Position); rows; cols }

let set_pipe field pos pipe = Hashtbl.add_exn field.pipes ~key:pos ~data:pipe
let get_pipe field pos = Hashtbl.find_exn field.pipes pos

exception Bad_pipe of char

let pipe_of_char = function
  | '|' -> Vert
  | '-' -> Horiz
  | 'L' -> NE
  | 'J' -> NW
  | '7' -> SW
  | 'F' -> SE
  | '.' -> Ground
  | 'S' -> Start
  | c -> raise (Bad_pipe c)

module Part1 = struct
  (* from current position, go one step in direction *)
  let go field Position.(Position (x, y)) =
    let open Position in
    function
    | N ->
        assert (x > 0);
        Position (x - 1, y)
    | E ->
        assert (y < field.cols - 1);
        Position (x, y + 1)
    | W ->
        assert (y > 0);
        Position (x, y - 1)
    | S ->
        assert (x < field.rows - 1);
        Position (x + 1, y)

  (* From current position, check for direction dir and pipe
     located there if we can go to the field. *)
  let viable dir pipe =
    match (dir, pipe) with
    | N, Vert | N, NW | N, NE -> true
    | E, Horiz | E, NE | E, SE -> true
    | S, Vert | S, SW | S, SE -> true
    | W, Horiz | W, NW | W, SW -> true
    | _, Start -> true
    | _ -> false

  (* would a step in direction dir onto a tile with pipe next be allowed? *)
  let valid_step dir next =
    match (dir, next) with
    | N, Vert | N, SW | N, SE -> true
    | E, Horiz | E, NW | E, SW -> true
    | S, Vert | S, NW | S, NE -> true
    | W, Horiz | W, NE | W, SE -> true
    | _, Start -> true
    | _ -> false

  (* which directions can we go, given the current pipe? *)
  let dir_candidates pipe =
    let dirs = [ N; E; S; W ] in
    let f d = viable d pipe in
    List.filter ~f dirs

  (* Which positions are reachable from the current pos given the field? *)
  let next_candidates field pos =
    let open Position in
    let inner current =
      (* directions we can go based on the current pipe *)
      let dir_cands = dir_candidates current in
      (* neighbors we can transfer to based on the current pipe *)
      let pos_cands =
        List.map ~f:(fun dir -> (dir, go field pos dir)) dir_cands
      in
      (* neighbors that we can transfer to with a fitting pipe section *)
      let good_neighbor_filter (dir, pos) =
        if valid_step dir (get_pipe field pos) then Some pos else None
      in
      let pos_cands_good = List.filter_map ~f:good_neighbor_filter pos_cands in
      assert (2 >= List.length pos_cands_good);
      pos_cands_good
    in
    match get_pipe field pos with current -> inner current | exception _ -> []

  (* Beginning at start, discover the path, then return the path once back at the start. *)
  let rec traverse ?(other = false) field start current path =
    let cands = next_candidates field current in
    let empty_path = List.is_empty path
    and at_start = Position.eq start current in
    let finished = (not empty_path) && at_start in
    if finished then path
    else
      match path with
      | [] ->
          (* start position *)
          assert (0 = Position.compare start current);
          assert (2 = List.length cands);
          (* start with which candidate? determines direction of traversal *)
          let ix = if other then 1 else 0 in
          let first = List.nth_exn cands ix in
          traverse ~other field start first (current :: path)
      | last :: _lasts ->
          let not_last p = not (0 = Position.compare p last) in
          let nexts = List.filter ~f:not_last cands in
          let next = List.hd_exn nexts in
          traverse ~other field start next (current :: path)

  (* Find start point and traverse path. *)
  let find_path field =
    let start = find_start field.pipes in
    traverse field start start []
end

module Parse = struct
  let parse_row str =
    let chars = String.to_list str in
    List.map ~f:pipe_of_char chars

  let to_field parsed_rows =
    let field =
      make_field (List.length parsed_rows)
        (List.length @@ List.hd_exn parsed_rows)
    in
    let insert_row r c p = set_pipe field (Position.Position (r, c)) p in
    let f r row = List.iteri ~f:(insert_row r) row in
    List.iteri ~f parsed_rows;
    field

  let parse_field rows =
    let parsed_rows = List.map ~f:parse_row rows in
    to_field parsed_rows
end

module Part2 = struct
  (* bitmap for tiles: true if tile is inside of path. *)
  type t = bool array

  let create field : t = Array.create ~len:(field.rows * field.cols) false

  (* check if position is valid *)
  let exists field (Position.Position (x, y)) =
    x < field.rows && y < field.cols && x >= 0 && y >= 0

  (* convert position to index in bitmap *)
  let pos_to_ix field (Position.Position (x, y)) = (field.cols * x) + y
  (* convert index (in bitmap) to position *)

  let ix_to_pos field ix =
    let row = Int.(ix / field.cols) and col = Int.(ix % field.cols) in
    Position.Position (row, col)

  (* check if pos is marked as inside the field. *)
  let is_inside field t pos =
    let ix = pos_to_ix field pos in
    if ix >= 0 && ix < Array.length t then t.(ix) else false

  (* mark a tile as inside of the path *)
  let mark_inside field t pos = t.(pos_to_ix field pos) <- true

  (* going from from to top, which direction are we going? *)
  let direction from top =
    let open Position in
    let diff a b = Int.(abs (a - b)) in
    let (Position (fx, fy)) = from and (Position (tx, ty)) = top in
    assert (diff fx tx + diff fy ty = 1);
    match (tx - fx, ty - fy) with
    | 1, 0 -> S
    | -1, 0 -> N
    | 0, 1 -> E
    | 0, -1 -> W
    | _ -> assert false

  (* return rel. positions to the left of pipe if we go into it heading direction dir.
     For example, a horizontal pipe entered from the east going west has its left element
     in the cell right above it.
  *)
  let left_of_pipe pipe dir =
    match (pipe, dir) with
    | Vert, N -> [ (0, -1) ]
    | Vert, S -> [ (0, 1) ]
    | Vert, _ -> assert false
    | Horiz, W -> [ (1, 0) ]
    | Horiz, E -> [ (-1, 0) ]
    | Horiz, _ -> assert false
    | NE, S -> [ (-1, 1) ]
    | NE, W -> [ (1, 0); (1, -1); (0, -1) ]
    | NE, _ -> assert false
    | NW, S -> [ (0, 1); (1, 1); (1, 0) ]
    | NW, E -> [ (-1, -1) ]
    | NW, _ -> assert false
    | SE, N -> [ (-1, -1); (0, -1); (-1, 0) ]
    | SE, W -> [ (1, 1) ]
    | SE, _ -> assert false
    | SW, N -> [ (1, -1) ]
    | SW, E -> [ (-1, 0); (-1, 1); (0, 1) ]
    | SW, _ -> assert false
    | Ground, _ -> assert false
    (* a start field is accessible from all directions *)
    | Start, d -> (
        match d with
        | N -> [ (0, -1) ]
        | S -> [ (0, 1) ]
        | E -> [ (-1, 0) ]
        | W -> [ (1, 0) ])

  (* on the segment from -> to_pos, generate all positions around to_pos
     that are to the left of it (in the direction of traversal. *)
  let candidate_positions from to_pos to_pipe =
    let open Position in
    let dir = direction from to_pos in
    let (Position (px, py)) = to_pos in
    let f (dx, dy) = Position (px + dx, py + dy) in
    List.map ~f (left_of_pipe to_pipe dir)

  (* for each segment of the path (from -> top), mark all left neighbors.
     A neighbor is marked if it is not part of the path. The tile type doesn't matter. *)
  let mark_left_neighbors field t pathset from top =
    let pipe = get_pipe field top in
    let cands = candidate_positions from top pipe in
    let not_in_path pos = not (Hash_set.mem pathset pos) in
    let neighs =
      List.filter ~f:(fun c -> exists field c && not_in_path c) cands
    in
    let mark pos = mark_inside field t pos in
    List.iter ~f:mark neighs

  (* mark ground tiles left-adjacent to path. *)
  let mark_initial field t path pathset =
    let f current next =
      mark_left_neighbors field t pathset current next;
      next
    in
    ignore @@ List.fold ~init:(List.hd_exn path) ~f (List.tl_exn path)

  (* return candidate positions neighboring pos. Not all of them are valid, for example they might be part of the path. *)
  let neighbor_cands pos =
    let open Position in
    let (Position (x, y)) = pos in
    let f (dx, dy) = Position (x + dx, y + dy) in
    let cands_d = [ (1, 0); (0, 1); (-1, 0); (0, -1) ] in
    List.map ~f cands_d

  (* mark any neighbors of pos if they are ground.
     No diagonal neighbors considered.
     This is basically DFS. *)
  let rec mark_adjacent_ground field t pathset pos =
    let neighs = neighbor_cands pos in
    let ok neigh =
      (not @@ Hash_set.mem pathset neigh)
      && (not (is_inside field t neigh))
      && exists field neigh
    in
    let ok_neighs = List.filter ~f:ok neighs in
    List.iter ~f:(mark_inside field t) ok_neighs;
    (* recurse into neighbors (DFS) *)
    List.iter ~f:(mark_adjacent_ground field t pathset) ok_neighs

  (* entry point: set up inside bitmap, mark initial left-of-path neighbors, and expand
     from neighbors to all reachable tiles within the path. *)
  let count_enclosed_tiles field path =
    let t = create field in
    (* in addition to the path list we supply a hash set of all nodes, which drastically
       accelerates the marking process. *)
    let pathset = Hash_set.of_list (module Position) path in
    mark_initial field t path pathset;
    let mark ix marked =
      if marked then mark_adjacent_ground field t pathset (ix_to_pos field ix)
      else ()
    in
    Array.iteri ~f:mark t;
    Array.count ~f:(fun x -> x) t
end

let () =
  let input = In_channel.(input_lines stdin) in
  let field = Parse.parse_field input in
  let path = Part1.find_path field in
  let revpath = List.rev path in
  (* dirty trick: we don't know which direction the traversal took.
     But the direction determines what's "inside" and what's "outside".
     Therefore: just calculate both inside and outside tiles, one of them
     will be right *)
  let in_path1 = Part2.count_enclosed_tiles field revpath in
  let in_path2 =
    Part2.count_enclosed_tiles field (List.hd_exn revpath :: path)
  in
  Out_channel.(
    printf "length: %d\n" (List.length path);
    printf "in path: %d %d\n" in_path1 in_path2)