:- module day16.
:- 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 map.
:- import_module digraph.
:- import_module set.

:- type label == string.
:- type key   == digraph_key(label).
:- type rates == map(key, int).
:- type graph == digraph(label).
:- type open  == set(key).


:- pred spaces(chars::in, chars::out) is det.
spaces --> if [' '] then spaces else [].

:- pred word0(chars::out, chars::in, chars::out) is semidet.
word0([C | Cs]) -->
  [C], {is_alpha(C)},
  (if word0(Cs0) then {Cs = Cs0} else {Cs = []}).

:- pred word(string::out, chars::in, chars::out) is semidet.
word(W) --> word0(Cs), {W = from_char_list(Cs)}, spaces.

:- pred words(list(string)::in, chars::in, chars::out) is semidet.
words([])       --> [].
words([W | Ws]) --> word(W), words(Ws).

:- pred sym(char::in, chars::in, chars::out) is semidet.
sym(S) --> [S], spaces.

:- pred tunnels(list(string)::out, chars::in, chars::out) is semidet.
tunnels([T | Ts]) -->
  word(T), (if sym(',') then tunnels(Ts) else {Ts = []}).

:- pred number(int::out, chars::in, chars::out) is semidet.
number(N) --> digits(S), {to_int(from_char_list(S), N)}.

:- pred digits(chars::out, chars::in, chars::out) is semidet.
digits([C | Cs]) -->
  [C], {is_digit(C)},
  (if digits(Cs2) then {Cs = Cs2} else {Cs = []}).

:- pred line(graph::in, graph::out, rates::in, rates::out,
             chars::in, chars::out) is nondet.
line(!Graph, !Rates) -->
  word("Valve"), word(Label),
  words(["has", "flow", "rate"]), sym('='),
  number(Rate), sym(';'),
  (words(["tunnels", "lead",  "to", "valves"]) ;
   words(["tunnel",  "leads", "to", "valve"])),
  tunnels(Tunnels),
  {add_vertex(Label, Key, !Graph)},
  {insert(Key, Rate, !Rates)},
  {foldl(add_vertices_and_edge(Label), Tunnels, !Graph)}.

:- pred line(string::in, graph::in, graph::out, rates::in, rates::out)
    is nondet.
line(Str, !G, !R) :- line(!G, !R, to_char_list(Str), []).

:- pred lines(lines::in, graph::out, rates::out) is cc_nondet.
lines(Lines, G, R) :- foldl2(line, Lines, digraph.init, G, map.init, R).


:- pred max_release(int::in, label::in, graph::in, rates::in, int::out) is det.
max_release(Time, Start, Graph, Rates, Max) :-
  max_release(Time, lookup_key(Graph, Start), Graph, Rates, init, _, Max).

:- pred max_release(int, key, graph, rates, open, open, int).
:- mode max_release(in, in, in, in, in, out, out) is det.
:- pragma memo(max_release/7,
    [specified([value, value, addr, addr, value, output, output])]).
max_release(Time, Here, Graph, Rates, !Open, Max) :-
  if Time = 0 then Max = 0 else
  HereRate = lookup(Rates, Here),
  StartOpen = !.Open,
  (if
    HereRate \= 0,
    insert_new(Here, !Open)
  then
    max_release(Time - 1, Here, Graph, Rates, !Open, Max0),
    MaxHere = Max0 + HereRate * (Time - 1)
  else
    MaxHere = 0),
  Nexts = [Here | set.to_sorted_list(lookup_from(Graph, Here))],
  max_releases(Time - 1, Nexts, Graph, Rates, StartOpen, Outs),
  get_max([o(!.Open, MaxHere) | Outs], o(!:Open, Max)).

:- type out ---> o(open :: open, max :: int).

:- pred max_releases(int, list(key), graph, rates, open, list(out)).
:- mode max_releases(in, in, in, in, in, out) is det.
max_releases(_, [], _, _, _, []).
max_releases(Time, [Here|Theres], Graph, Rates, Open0, [o(Open, Max) | Rest]) :-
  max_release(Time, Here, Graph, Rates, Open0, Open, Max) &
  max_releases(Time, Theres, Graph, Rates, Open0, Rest).


:- import_module bool.

:- func tick(bool, int) = int.
tick(yes, N) = N - 1.
tick(no,  N) = N.


:- pred max_release2(int::in, label::in, graph::in, rates::in, int::out) is det.
max_release2(Time, Start, Graph, Rates, Max) :-
  die("this works on the example but it's too inefficient for the real thing"),
  Key = lookup_key(Graph, Start),
  max_release2(Time, no, Key, Key, Graph, Rates, init, _, Max).

:- pred max_release2(int, bool, key, key, graph, rates, open, open, int).
:- mode max_release2(in, in, in, in, in, in, in, out, out) is det.
:- pragma memo(max_release2/9,
    [specified([value, value, value, value, addr, addr,
                value, output, output])]).
max_release2(Time, Tick, Here1, Here2, Graph, Rates, !Open, Max) :-
  if Time = 0 then Max = 0 else
  HereRate = lookup(Rates, Here1),
  StartOpen = !.Open,
  (if
    HereRate \= 0,
    insert_new(Here1, !Open)
  then
    max_release2(tick(Tick, Time), not(Tick), Here2, Here1,
                 Graph, Rates, !Open, Max0),
    MaxHere = Max0 + HereRate * (Time - 1)
  else
    MaxHere = 0),
  Nexts = [Here1 | set.to_sorted_list(lookup_from(Graph, Here1))],
  max_releases2(tick(Tick, Time), not(Tick), Nexts, Here2,
                Graph, Rates, StartOpen, Outs),
  get_max([o(!.Open, MaxHere) | Outs], o(!:Open, Max)).

:- pred max_releases2(int, bool, list(key), key, graph, rates, open, list(out)).
:- mode max_releases2(in, in, in, in, in, in, in, out) is det.
max_releases2(_, _, [], _, _, _, _, []).
max_releases2(Time, Tick, [Here1|Theres], Here2, Graph, Rates, Open0,
              [o(Open, Max) | Rest]) :-
  max_release2(Time, Tick, Here2, Here1, Graph, Rates, Open0, Open, Max) &
  max_releases2(Time, Tick, Theres, Here2, Graph, Rates, Open0, Rest).


:- pred get_max(list(out)::in(non_empty_list), out::out) is det.
get_max([O], O).
get_max([O, R | Rest], O0) :-
  get_max([R | Rest], OR),
  (if O^max >= OR^max then O0 = O else O0 = OR).

run(Part, Lines, int(Max)) :-
  if lines(Lines, G, R) then
    (if Part = one then max_release(30, "AA", G, R, Max)
    else max_release2(26, "AA", G, R, Max))
  else die("bad input").