:- module day11.
:- interface.
:- import_module basics.

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

:- implementation.
:- import_module int.
:- import_module string.
:- import_module list.
:- import_module array.

:- type item == int.

:- type op ---> add(int); mul(int); square.
:- type test == int.

:- type monkey --->
    m(items :: list(item),
      op :: op, test :: test,
      if_true :: int, if_false :: int,
      inspected :: int).
:- func inspected(monkey) = int.
:- func test(monkey) = int.

:- type monkeys == array(monkey).
:- mode m_in  == array_di.
:- mode m_out == array_uo.

:- type divmod ---> dm(d :: int, m :: int).

:- func op(op, item) = item.
op(add(M), N) = M + N.
op(mul(M), N) = M * N.
op(square, N) = N * N.

:- func op(op, divmod, item) = item.
op(O, dm(D, M), N) = (op(O, N) div D) mod M.

:- pred test(test::in, item::in) is semidet.
test(M, N) :- N mod M = 0.


:- import_module parsing_utils.

eof(Src) --> not next_char(Src, _).

:- pred kw(src::in, string::in, ps::in, ps::out) is semidet.
kw(Src, Kw) --> ikeyword("abcdefghijklmnopqrstuvwxyz", Kw, Src, _).

:- pred kws(src::in, list(string)::in, ps::in, ps::out) is semidet.
kws(_,   [])         --> [].
kws(Src, [Kw | Kws]) --> kw(Src, Kw), kws(Src, Kws).

colon(Src) --> punct(":", Src, _).
comma(Src) --> punct(",", Src, _).
equal(Src) --> punct("=", Src, _).
plus(Src) --> punct("+", Src, _).
star(Src) --> punct("*", Src, _).


:- pred monkeys(src::in, list(monkey)::out, ps::in, ps::out) is semidet.
monkeys(Src, Ms) --> one_or_more(monkey, Src, Ms), eof(Src).

:- pred monkey(src::in, monkey::out, ps::in, ps::out) is semidet.
monkey(Src, m(reverse(I), O, P, T, F, 0)) -->
  kw(Src, "monkey"), int_literal_as_string(Src, _), colon(Src),
  kws(Src, ["starting", "items"]), colon(Src), items(Src, I),
  kw(Src, "operation"), colon(Src),
    kw(Src, "new"), equal(Src), op(Src, O),
  kw(Src, "test"), colon(Src), test(Src, P),
  kws(Src, ["if", "true"]), colon(Src), act(Src, T),
  kws(Src, ["if", "false"]), colon(Src), act(Src, F).

:- pred items(src::in, list(item)::out, ps::in, ps::out) is semidet.
items(Src, Is) --> comma_separated_list(int_literal, Src, Is).

:- pred op(src::in, op::out, ps::in, ps::out) is semidet.
op(Src, Op) -->
  kw(Src, "old"),
  (if star(Src) then
    (if int_literal(Src, N) then {Op = mul(N)}
     else kw(Src, "old"), {Op = square})
   else plus(Src), int_literal(Src, N), {Op = add(N)}).

:- pred test(src::in, test::out, ps::in, ps::out) is semidet.
test(Src, N) --> kws(Src, ["divisible", "by"]), int_literal(Src, N).

:- pred act(src::in, int::out, ps::in, ps::out) is semidet.
act(Src, N) --> kws(Src, ["throw", "to", "monkey"]), int_literal(Src, N).


:- pred give(item::in, int::in, monkeys::m_in, monkeys::m_out) is det.
give(I, N, !Ms) :-
  lookup(!.Ms, N, M0),
  M = M0^items := [I | M0^items],
  set(N, M, !Ms).

:- pred get(int::in, monkey::out, monkeys::m_in, monkeys::m_out) is det.
get(N, M0, !Ms) :-
  lookup(!.Ms, N, M0),
  Len = length(M0^items), Insp = M0^inspected,
  M = (M0^inspected := Insp + Len)^items := [],
  set(N, M, !Ms).

:- pred step(divmod::in, monkey::in, item::in,
             monkeys::m_in, monkeys::m_out) is det.
step(DM, M, I, !Ms) :-
  I2 = op(M^op, DM, I),
  (if test(M^test, I2) then
    give(I2, M^if_true, !Ms)
  else
    give(I2, M^if_false, !Ms)).

:- pred steps(divmod::in, monkey::in, list(item)::in,
              monkeys::m_in, monkeys::m_out) is det.
steps(_,  _, [],       !Ms).
steps(DM, M, [I | Is], !Ms) :- steps(DM, M, Is, !Ms), step(DM, M, I, !Ms).

:- pred turn(divmod::in, int::in, monkeys::m_in, monkeys::m_out) is det.
turn(DM, N, !Ms) :- get(N, M, !Ms), steps(DM, M, M^items, !Ms).

:- pred round(divmod::in, int::in, monkeys::m_in, monkeys::m_out) is det.
round(DM, N, !Ms) :-
  if N < size(!.Ms) then turn(DM, N, !Ms), round(DM, N + 1, !Ms) else true.

:- pred round(divmod::in, monkeys::m_in, monkeys::m_out) is det.
round(DM, !Ms) :- round(DM, 0, !Ms).

:- pred rounds(divmod::in, int::in, monkeys::m_in, monkeys::m_out) is det.
rounds(DM, N, !Ms) :-
  if N > 0 then round(DM, !Ms), rounds(DM, N - 1, !Ms) else true.

:- pred rounds(part::in, monkeys::m_in, monkeys::m_out) is det.
rounds(Part, !Ms) :-
  DM = dm(part(Part, 3, 1), foldl(func(X, A) = test(X) * A, !.Ms, 1)),
  rounds(DM, part(Part, 20, 10_000), !Ms).


% i broke part 1 at some point and i can't see where

run(Part, Lines, Out) :-
  parse(join_list("\n", Lines), monkeys, Res),
  (if Res = ok(MList) then
    from_list(MList, Arr0),
    rounds(Part, Arr0, Arr1),
    to_list(Arr1, Final),
    Out = int(prod(take_upto(2, rev_sort(map(inspected, Final)))))
  else
    Out = 'new other'(Res)).