:- module day13.
:- 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.

:- type packet ---> i(int); l(list(packet)).
:- type packets == {packet, packet}.


:- pred packets(list(packets)::out, lines::in) is nondet.
packets(Ps, Lines) :- packets(Ps, to_char_list(join_list("", Lines)), []).

:- pred packets(list(packets)::out, chars::in, chars::out) is multi.
packets([{P, Q} | Rest]) --> packet(P), packet(Q), packets(Rest).
packets([])              --> [].

:- pred packet(packet::out, chars::in, chars::out) is nondet.
:- import_module day4.
packet(i(I))  --> day4.number(I).
packet(l(Ps)) --> ['['], inner(Ps), [']'].

:- pred inner(list(packet)::out, chars::in, chars::out) is multi.
inner(Ps) --> inner1(Ps).
inner([]) --> [].

:- pred inner1(list(packet)::out, chars::in, chars::out) is nondet.
inner1([P | Ps]) --> packet(P), (if [','] then inner1(Ps) else {Ps = []}).


:- pred cmps(list(packet)::in, list(packet)::in, comparison_result::out) is det.
cmps([],     [],     =).
cmps([],     [_|_],  <).
cmps([_|_],  [],     >).
cmps([P|Ps], [Q|Qs], C) :-
  cmp(P, Q, C1), (if C1 = (=) then cmps(Ps, Qs, C) else C = C1).

:- pred cmp(packet::in, packet::in, comparison_result::out) is det.
cmp(i(I),  i(J),  C) :- compare(C, I, J).
cmp(l(Ps), l(Qs), C) :- cmps(Ps,     Qs,     C).
cmp(i(I),  l(Qs), C) :- cmps([i(I)], Qs,     C).
cmp(l(Ps), i(J),  C) :- cmps(Ps,     [i(J)], C).

:- func cmp(packet, packet) = comparison_result.
cmp(P, Q) = C :- cmp(P, Q, C).

:- pred ok(packets::in) is semidet.
ok({P, Q}) :- not cmp(P, Q, >).

div1 = l([l([i(2)])]).
div2 = l([l([i(6)])]).

:- pred divider(packet::in) is semidet.
divider(div1). divider(div2).

:- mode unary == (pred(in) is semidet).

:- pred indices(pred(T)::unary, int::in, list(T)::in, list(int)::out) is det.
indices(_, _, [],     []).
indices(P, I, [X|Xs], Is) :-
  indices(P, I + 1, Xs, Is0),
  (if P(X) then Is = [I|Is0] else Is = Is0).

:- pred indices(pred(T)::unary, list(T)::in, list(int)::out) is det.
indices(P, Xs, Is) :- indices(P, 1, Xs, Is).


:- func merge(list(packets)) = list(packet).
merge(Ps) = foldl(func({X, Y}, Zs) = [X,Y|Zs], Ps, []).

:- pragma no_determinism_warning(run/3).
run(one, Lines, int(sum(Is))) :-
  if packets(Ps, Lines) then
    indices(ok, Ps, Is)
  else
    die("bad input").
run(two, Lines, int(prod(Is))) :-
  if packets(Pairs, Lines) then
    indices(divider, sort(cmp, [div1, div2 | merge(Pairs)]), Is)
  else
    die("bad input").