:- use_module(library(pio)).
:- initialization(main, main).

main([FileName | _]) :-
    input(FileName, Map),
    findall(Dist, distance(_, Map, _, Dist), Dists),
    sum_list(Dists, Sum),
    writef('Answer=%t\n', [Sum]).

% Dist is distance along Axis in Map between two row/col A & B (after expansion)
distance(Axis, Map, A-B, Dist) :-
    true_coords(Axis, Map, TrueNs),
    member(A-ACount, TrueNs), member(B-BCount, TrueNs), A < B,
    Dist is (B - A) * ACount * BCount.

% TrueXs is the list of all true coordinates along Axis in Map
true_coords(x, Map, TrueXs) :-
    NextX = [Row, LastX-_, X-Count]>>(
        sum_list(Row, Count),
        (Count = 0 -> X is LastX + 2; X is LastX + 1)),
    scanl(NextX, Map, -1-0, Xs), Xs = [_ | TrueXs].
true_coords(y, [FirstRow | Map], TrueYs) :-
    foldl(add_vectors, Map, FirstRow, CountY),
    NextY = [Count, LastY-_, Y-Count]>>(
        Count = 0 -> Y is LastY + 2; Y is LastY + 1),
    scanl(NextY, CountY, -1-0, Ys), Ys = [_ | TrueYs].

add_vectors([], [], []).
add_vectors([A | V1], [B | V2], [C | V]) :- C is A + B, add_vectors(V1, V2, V).

% Read file into 2D array Map. 1 corresponding to a galaxy and 0 otherwise.
input(FileName, Map) :- phrase_from_file(lines(Map), FileName).
lines([]) --> eos, !.
lines([Line|Lines]) --> line(Line), lines(Lines).
line([]) --> ( "\n" ; eos ), !.
line([L|Ls]) --> ([C], {C =:= 35 -> L = 1; L = 0}), line(Ls).
eos([], []).