:- module day14.
:- interface.
:- import_module basics.

:- pred run(part::in, lines::in, answer::out) is cc_multi.

:- implementation.
:- import_module int.
:- import_module char.
:- import_module string.
:- import_module list.
:- import_module array2d.
:- import_module ranges.

:- import_module set_bbbtree.
:- type set(T) == set_bbbtree(T).


:- type cell ---> empty; rock; sand.

:- type data == array2d(cell).

:- type grid --->
    g(start_x :: int, end_x :: int,
      start_y :: int, end_y :: int,
      data    :: data).

:- inst grid for grid/0 ---> g(ground, ground, ground, ground, array2d).
:- mode grid_di == di(grid).
:- mode grid_uo == out(grid).

:- type point == {int, int}.

holeX = 500. holeY = 0. hole = {holeX, holeY}.


:- use_module day4.
:- pred line(list(point)::out, chars::in, chars::out) is semidet.
line([{X, Y} | Rest]) -->
  day4.number(X), [','], day4.number(Y),
  (if [' ','-','>',' '] then line(Rest) else {Rest = []}).

:- pred line(string::in, list(point)::out) is semidet.
line(Line, Points) :- line(Points, to_char_list(Line), []).

:- pred lines(lines::in, list(list(point))::out) is semidet.
lines(Lines, Points) :- map(line, Lines, Points).


:- func expand(list(list(point))) = set(point).
expand(Points) = from_list(condense(map(expand1, Points))).

:- func expand1(list(point)) = list(point).
expand1([])            = [].
expand1([P])           = [P].
expand1([P, Q | Rest]) = expand(P, Q) ++ expand1([Q | Rest]).

:- func expand(point, point) = list(point).
expand({X1, Y1}, {X2, Y2}) = List :-
  if X1 = X2 then List = map(func(Y) = {X1, Y}, between(Y1, Y2))
  else if Y1 = Y2 then List = map(func(X) = {X, Y1}, between(X1, X2))
  else die("expand/2: non-orthogonal line").

:- func between(int, int) = list(int).
between(I, J) = to_sorted_list(range(min(I, J), max(I, J))).


:- func bounds(set(point)) = {point, point}.
bounds(Points) = {{LoX, LoY}, {HiX, HiY}} :-
  Points0 = insert(Points, hole),
  Xs = to_sorted_list(map(func({X, _}) = X, Points0)),
  Ys = to_sorted_list(map(func({_, Y}) = Y, Points0)),
  LoX = det_head(Xs), det_last(Xs, HiX),
  LoY = det_head(Ys), det_last(Ys, HiY).


:- func add_floor(set(point)) = set(point).
add_floor(Points) = insert_list(Points, Floor) :-
  {_, {_, Lo}} = bounds(Points),
  FloorY = Lo + 2, Start = holeX - FloorY, End = holeX + FloorY,
  Floor = map(func(X) = {X, FloorY}, between(Start, End)).


:- func get(grid, point) = cell.
get(G, {X, Y}) = G^data^elem(Y - G^start_y, X - G^start_x).

:- pred set_d(point::in, point::in, cell::in,
              data::array2d_di, data::array2d_uo) is det.
set_d({StartX, StartY}, {X, Y}, A, !Data) :-
  !Data^elem(Y - StartY, X - StartX) := A.

:- pred set(point::in, cell::in, grid::grid_di, grid::grid_uo) is det.
set(XY, A, !Grid) :- some [!Data] (
  !:Data = !.Grid^data,
  set_d({!.Grid^start_x, !.Grid^start_y}, XY, A, !Data),
  !Grid^data := !.Data
).

:- pred in_bounds(grid::in, point::in) is semidet.
in_bounds(G, {X, Y}) :-
  G^start_x =< X, X =< G^end_x,
  G^start_y =< Y, Y =< G^end_y.

:- pred place_rocks(point::in, list(point)::in,
                    data::array2d_di, data::array2d_uo) is det.
place_rocks(_, [], !Data).
place_rocks(Start, [XY | Rest], !Data) :-
  set_d(Start, XY, rock, !Data),
  place_rocks(Start, Rest, !Data).

:- func grid_from_set(set(point)) = grid.
:- mode grid_from_set(in) = grid_uo.
grid_from_set(Points) = G :-
  {{StartX, StartY}, {EndX, EndY}} = bounds(Points),
  Width = EndX - StartX + 1, Height = EndY - StartY + 1,
  place_rocks({StartX, StartY}, to_sorted_list(Points),
              init(Height, Width, empty), Arr),
  G = g(StartX, EndX, StartY, EndY, Arr).

:- func make_grid(lines) = grid.
:- mode make_grid(in) = grid_uo is semidet.
make_grid(Lines) = grid_from_set(expand(Points)) :-
  lines(Lines, Points).

:- func make_floor_grid(lines) = grid.
:- mode make_floor_grid(in) = grid_uo is semidet.
make_floor_grid(Lines) = grid_from_set(add_floor(expand(Points))) :-
  lines(Lines, Points).


:- type fall ---> stop(point); overflow.

:- pred move(point::in, fall::out, grid::grid_di, grid::grid_uo) is semidet.
move({X, Y}, B, !Grid) :-
  if not in_bounds(!.Grid, {X, Y}) then B = overflow
  else if get(!.Grid, {X, Y}) \= empty then fail
  else if move({X,   Y+1}, B0, !Grid) then B = B0
  else if move({X-1, Y+1}, B0, !Grid) then B = B0
  else if move({X+1, Y+1}, B0, !Grid) then B = B0
  else set({X, Y}, sand, !Grid), B = stop({X, Y}).

:- pred fill(int::out, grid::grid_di, grid::grid_uo) is det.
fill(N, !Grid) :-
  if move(hole, stop(_), !Grid) then fill(N0, !Grid), N = N0 + 1
  else N = 0.

:- pragma no_determinism_warning(run/3).
run(one, Lines, int(Res)) :-
  if Grid = make_grid(Lines) then fill(Res, Grid, _)
  else die("bad input").
run(two, Lines, int(Res)) :-
  if Grid = make_floor_grid(Lines) then fill(Res, Grid, _)
  else die("bad input").