:- lib(lists).
:- lib(util).
:- lib(ic).
:- lib(ic_global).
:- lib(ic_sets).
:- lib(branch_and_bound).
:- use_module(ai_util).

%%% Models %%%

%% N-queens %%
n_queens(N, Q) :-
	dim(Q, [N]),
	Q #:: 1..N,
	ic_global:alldifferent(Q),
	(for(I, 1, N), param(Q, N) do
		(for(J, I + 1, N), param(Q, I) do
			abs(Q[I] - Q[J]) #\= J - I,
			abs(Q[J] - Q[I]) #\= J - I
		)
	).

print_queens(Q, Chessboard) :-
	dim(Q, [N]),
	dim(Chessboard, [N, N]),
	(for(J, 1, N), param(Chessboard, Q, N) do
		(for(I, 1, N), param(Chessboard, Q, J) do
			I is Q[J] -> 1 is Chessboard[I, J]; 0 is Chessboard[I, J]
		)
	),
	print_term_matrix(Chessboard).

print_matrix(Matrix) :-
	dim(Matrix, [Rows, Cols]),
	(multifor([I,J], [1, 1], [Rows, Cols]), param(Matrix, Cols) do
		X is Matrix[I,J], printf('%3d', X), (J = Cols -> nl; true)
	).

%% Magic Sequence %%
magic_sequence(N, Seq) :-
	dim(Seq, [N]),
	Seq #:: 0..(N - 1),
	(for(I, 1, N), param(Seq) do
		X is Seq[I],
		I1 is I - 1,
		occurrences(I1, Seq, X)
	),
	(for(I, 1, N - 1), foreach(E, Expr), param(Seq) do
		E = I * Seq[I + 1]
	),
	sum(Expr) #= N.

%% Futoshiki %%
make_futoshiki(N, Hints, Constr, Matrix) :-
	dim(Matrix, [N, N]),
	Matrix #:: 1..N,
	(foreacharg(Row, Matrix) do ic:alldifferent(Row)),
	(for(J, 1, N), param(Matrix, N) do Col is Matrix[1..N, J], ic:alldifferent(Col)),
	(foreach(I-J > X, Hints), param(Matrix) do Matrix[I, J] #= X),
	(foreach(I1-J1 > I2-J2, Constr), param(Matrix) do Matrix[I1, J1] #> Matrix[I2, J2]).

solve_futoshiki(Matrix) :-
	labeling(Matrix),
	print_matrix(Matrix).

futoshiki_5_1(M) :- make_futoshiki(5
	, [3-2 > 4, 4-2 > 5, 5-1 > 5, 4-5 > 2]
	, [1-2 > 1-3, 2-1 > 1-1, 2-5 > 3-5, 4-4 > 3-4, 5-2 > 5-3]
	, M).

futoshiki_7_1(M) :- make_futoshiki(6
	, [4-3 > 2, 5-1 > 3, 5-4 > 6, 5-6 > 1, 6-5 > 6, 6-6 > 3]
	, [1-3 > 1-2, 2-3 > 1-3, 2-2 > 2-3, 3-2 > 4-2, 3-4 > 3-5, 3-5 > 4-5, 5-2 > 5-3]
	, M).

%% Magic Square %%
% slow for N > 6
magic_square(N, A) :-
	Sum is N *(N ^ 2 + 1) // 2, % magic constant
	dim(A, [N, N]), A #:: 1..N * N, % setup variables
	(for(I, 1, N), foreach(X, D1), foreach(Y, D2), param(A, Sum, N) do
		X is A[I, I], % fill primary diagonal
		Y is A[I, N - I + 1], % fill secondary diagonal
		sum(A[I, 1..N]) #= Sum, % row constraint
		sum(A[1..N, I]) #= Sum % column constraint
	),
	sum(D1) #= Sum, sum(D2) #= Sum, % diagonal constraints
	ic_global:alldifferent(A).

%% Latin Square %%
latin_square(N, Matrix) :-
	dim(Matrix, [N, N]),
	Matrix #:: 1..N,
	(for(I, 1, N), param(Matrix, N) do
		Row is Matrix[I, 1..N],
		Col is Matrix[1..N, I],
		ic_global:alldifferent(Row),
		ic_global:alldifferent(Col)
	).

%%% Golomb Rulers %%%
golomb_ruler(N, Seq, Length) :-
	dim(Seq, [N]),
	Seq #:: 0..1.0Inf,
	Seq[1] #= 0,
	Length #= Seq[N],
	ordered(<, Seq),
	(for(I, 1, N), foreach(Xs, Ds), param(Seq, N) do
		(for(J, I + 1, N), foreach(X, Xs), param(Seq, I) do
			X #= Seq[J] - Seq[I]
		)
	),
	flatten(Ds, Differences),
	ic:alldifferent(Differences),
	append([D1|_], [Dn], Differences), !,
	D1 #< Dn.

%% All-intervals %%
all_intervals(N, Seq) :-
	dim(Seq, [N]),
	Seq #:: 0..N - 1,
	ic:alldifferent(Seq),
	Begin is Seq[1..N - 1], End is Seq[2..N],
	(foreach(B, Begin), foreach(E, End), foreach(D, Ds) do
		D #= abs(E - B)
	),
	Ds #:: 1..N - 1,
	ic:alldifferent(Ds).

%% Map Coloring %%
:- local struct(country(name, color)).

countries([wa, nt, q, nsw, v, sa, t]).

n(wa, nt).
n(nt, q).
n(q, nsw).
n(nsw, v).
n(sa, wa).
n(sa, nt).
n(sa, q).
n(sa, nsw).
n(sa, v).

map_coloring(Countries, MaxColors) :-
	countries(CountryNames),
	(foreach(Name, CountryNames), foreach(country{name:Name, color:V}, Countries), param(MaxColors) do
		V #:: 1..MaxColors
	),
	findall(n(A,B), n(A,B), L),
	(foreach(n(A, B), L), param(Countries) do
		member(country{name:A, color:Va}, Countries),
		member(country{name:B, color:Vb}, Countries),
		Va #\= Vb
	),
	search(Countries,
	       color of country, % select based on this variable
	       most_constrained,
	       indomain,
	       complete,
	       [backtrack(Back)]),
	writeln(Countries),
	writeln(Back).

%% Graph Coloring %%
graph_coloring(Adj, Colors, Vertices) :-
	dim(Adj, [N, N]),
	dim(Vertices, [N]), Vertices #:: 1..Colors,
	Vertices[1] #= 1, % some symmetry break
	(for(J, 2, N), fromto(L, In, Out, _), param(Adj, Vertices) do
		1 is Adj[1, J] -> V is Vertices[J], In = [V|Out]; In = Out
	),
	ordered(<, L),
	(for(I, 1, N), param(N, Adj, Vertices) do
		(for(J, I + 1, N), param(I, Adj, Vertices) do
			Adj[I, J] => Vertices[I] #\= Vertices[J]
		)
	).

% Sudoku-like formulation %
graph_coloring_2(Adj, Colors, Vertices) :-
	dim(Adj, [N, N]),
	dim(Vertices, [N]), Vertices #:: 1..Colors,
	Vertices[1] #= 1,
	(foreacharg(Row, Adj), foreachelem(V, Vertices), param(Vertices) do
		(foreachelem(A, Row), foreachelem(N, Vertices), fromto(Neighbors, In, Out, []) do
			A = 1 -> In = [N|Out]; In = Out
		),
		ic:alldifferent([V|Neighbors])
	).

% Edge-list formulation %
graph_coloring_el(N, Edges, Vertices, Colors) :-
	dim(Vertices, [N]),
	Vertices #:: 1..N,
	(foreach(e(A, B), Edges), param(Vertices) do
		Vertices[A] #\= Vertices[B]
	),
	max(Vertices) #= Colors.

%% Clique %%
clique(Adj, Vertices, M) :-
	dim(Adj, [N, N]),
	dim(Vertices, [N]),
	Vertices #:: 0..1,
	flatten_array(Vertices, VerticesList),
	sum(VerticesList) #= M,
	(foreacharg(Row, Adj), foreachelem(V, Vertices), param(Vertices, M) do
		(foreachelem(E, Row, [J]), fromto(Neigh, In, Out, []), param(Vertices) do
			E = 1 -> (X is Vertices[J], In = [X|Out]); In = Out
		),
		V => sum(Neigh) + 1 #= M
	).

%% a_list %%
a_list(List, L) :-
	L #:: 1..100,
	indomain(L, max),
	length(List, L),
	List #:: 0..(L - 1),
	ic:alldifferent(List),
	(fromto(List, [X, Y|T], [Y|T], [_]) do
		2 * X #< Y
	).

%% Steiner Triplets %%
steiner_triplets(N, Triplets) :-
	M is N * (N - 1) // 6,
	intsets(Triplets, M, 1, N),
	(foreach(T, Triplets) do
		#(T, 3)
	),
	(fromto(Triplets, [H|T], T, []) do
		(foreach(X, T), param(H) do
			#(H /\ X) #=< 1
		)
	).

set_label(S) :-
	(foreach(E, S) do
		insetdomain(E,_,_,_)
	).

all_subsets([], [[]]) :- !.
all_subsets([H|T], L) :-
	all_subsets(T, S),
	(foreach(X, S), fromto(L, [X, [H|X]|Tx], Tx, []), param(H) do true).

set_partitions(S, N, Min, Max) :-
	intsets(S, N, Min, Max),
	all_disjoint(S),
	findall(X, between(Min, Max, X), R),
	all_union(S, R).

%% Shur's Lemma %%
all_triplets(N, Ts) :-
	findall(X, between(1, N, X), R),
	all_subsets(R, Rsub),
	(foreach(X, Rsub), fromto(Ts, In, Out, []) do
		(X = [A,B,C], C is A + B)-> In = [X|Out]; In = Out
	).

shur_lemma(N, K, S) :-
	set_partitions(S, K, 1, N),
	all_triplets(N, Ts),
	(foreach(T, Ts), param(S) do
		(foreach(X, S), param(T) do
			#(X /\ T) #< 3
		)
	).

%% Mine Sweeper %%
mine_sweeper(Hints, Field, Mines) :-
	dim(Hints, [M, N]),
	dim(Field, [M, N]),
	(foreachelem(Cell, Field, [I, J]), foreachelem(Hint, Hints), param(Field) do
		integer(Hint) ->
			neighbors(I, J, Field, Neighbors),
			Cell #= 0,
			sum(Neighbors) #= Hint;
			Cell #:: 0..1
	),
	flatten_array(Field, Cells),
	sum(Cells) #= Mines.

input_mine_sweeper(File, NMines, Hints) :-
	read_lines(File, [Mines|Field]),
	number_string(NMines, Mines),
	eclipse_feat:read_matrix(Field, Hints).

print_minefield(Hints) :-
	(foreacharg(Row, Hints) do
		(foreachelem(Cell, Row) do
			(var(Cell) -> write('_'); write(Cell)),
			write('  ')
		),
		nl
	).

print_minefield_cell(Cell, FieldCell) :-
	var(Cell), !,
	(FieldCell = 1 -> write('X'); write('_')).

print_minefield_cell(Cell, _) :-
	write(Cell).

print_minefield(Hints, Field) :-
	(foreacharg(Row, Hints), foreacharg(FieldRow, Field) do
		(foreachelem(Cell, Row), foreachelem(FieldCell, FieldRow) do
			print_minefield_cell(Cell, FieldCell),
			write('  ')
		),
		nl
	).

neighbors(I, J, Matrix, Neighbors) :-
	dim(Matrix, [M, N]),
	Imin is max(1, I - 1), Imax is min(M, I + 1),
	Jmin is max(1, J - 1), Jmax is min(N, J + 1),
	Slice is Matrix[Imin..Imax, Jmin..Jmax],
	flatten(Slice, Neighbors).

%% Dominator Number %%
queen_attack(I, J, Row, Col) :-
	I =:= Row; J =:= Col; Row - Col =:= I - J; Row + Col =:= I + J.

bishop_attack(I, J, Row, Col) :-
	Row - Col =:= I - J; Row + Col =:= I + J.

king_attack(I, J, Row, Col) :-
	X is abs(I - Row), Y is abs(J - Col), X < 2, Y < 2.

covered(Row, Col, Chessboard) :-
	(foreachelem(Cell, Chessboard, [I, J]),
		fromto([], In, Out, L), param(Row, Col) do
		queen_attack(I, J, Row, Col) -> Out = [Cell|In]; In = Out
	),
	sum(L) #> 0.

domination(N, Chessboard, S) :-
	dim(Chessboard, [N, N]),
	Chessboard #:: 0..1,
	(multifor([I, J], 1, N), param(Chessboard) do
		covered(I, J, Chessboard)
	),
	term_variables(Chessboard, Vars),
	sum(Vars) #= S.

%%% Solvers %%%

solve(n_queens(N, Queens)) :-
	n_queens(N, Queens),
	search(Queens, 0, input_order, indomain_middle, complete, [backtrack(B)]),
	print_queens(Queens, _),
	printf('Backtracks: %d\n', B).

solve(latin_square(N, Matrix)) :-
	latin_square(N, Matrix),
	search(Matrix, 0, first_fail, indomain, complete, [backtrack(B)]),
	(foreacharg(Row, Matrix) do
		(foreachelem(E, Row) do write(E), write(' ')),
		nl
	),
	printf('Backtracks: %d\n', [B]).

solve(magic_sequence(N, Seq)) :-
	magic_sequence(N, Seq),
	search(Seq, 0, first_fail, indomain, complete, [backtrack(B)]),
	printf('Backtracks: %d\n', [B]).

solve(golomb_ruler(N, Seq)) :-
	golomb_ruler(N, Seq, Length),
	minimize(search(Seq, 0, input_order,
		% indomain, % this doesn't work (infinite loop)
		indomain_min, % this works
		complete, [backtrack(B)]), Length),
	printf('Backtracks: %d\n', [B]).

solve(all_intervals(N, Seq)) :-
	all_intervals(N, Seq),
	labeling(Seq).