:- use_module(library(bounds)).
:- use_module(library(assoc)).

all_diff([]).
all_diff([X|Xs]) :-
	diff_from_all(Xs, X),
	all_diff(Xs).

diff_from_all([], _).
diff_from_all([A|As], D) :-
	dif(A, D),
	diff_from_all(As, D).

seq(N, N).
seq(N0, N) :-
	succ(N0, N1),
	seq(N1, N).

max_num_element(Multiset, MaxNum) :-
	sort(Multiset, Set),
	empty_assoc(E),
	create_assoc(Set, E, Assoc0),
	collect_stats(Multiset, Assoc0, Assoc),
	assoc_to_list(Assoc, Stats0),
	maplist(swap, Stats0, Stats1),
	keysort(Stats1, Stats2),
	last(Stats2, MaxNum-_).

swap(A-B, B-A).
	
collect_stats([], A, A).
collect_stats([E|Es], A0, A) :-
	get_assoc(E, A0, Val0),
	succ(Val0, Val1),
	put_assoc(E, A0, Val1, A1),
	collect_stats(Es, A1, A).
	
create_assoc([], A, A).
create_assoc([E|Es], A0, A) :-
	put_assoc(E, A0, 0, A1),
	create_assoc(Es, A1, A).

	
ukodus(Multiset, Squares) :-
	(  Multiset == []
	-> Grid = []
	;  length(Multiset, Num),
	   max_num_element(Multiset, MaxNum),
	   seq(Num, Size),
	   [X,Y] in 1..Size,
	   X #>= MaxNum,
	   Y #>= MaxNum,
	   Y #=< X,
	   X * Y #= Size,
	   label([X,Y]),
	   length(Grid, X),
	   maplist(length_(Y), Grid),
	   maplist(all_diff, Grid),
	   cols_diff(Grid),
	   flatten(Grid, Vars),
	   assign(Multiset, Vars),
	   grid_to_squares(Grid, 0, Squares, [])
	).

grid_to_squares([], _) --> [].
grid_to_squares([Row|Rows], Col) -->
	grid_to_squares(Row, 0, Col),
	{ succ(Col, Col1) },
	grid_to_squares(Rows, Col1).

grid_to_squares([], _, _) --> [].
grid_to_squares([V|Vs], Row, Col) -->
	{ var(V) },
	!,
	{ succ(Row, Row1) },
	grid_to_squares(Vs, Row1, Col).
grid_to_squares([V|Vs], Row, Col) -->
	[square(Row, Col, V)],
	{ succ(Row, Row1) },
	grid_to_squares(Vs, Row1, Col).
	

assign([], _).
assign([M|Ms], Vars0) :-
	select(M, Vars0, Vars1),
	assign(Ms, Vars1).
	
length_(X, Y) :- length(Y, X).

cols_diff(Rows) :-
	Rows = [R|_],
	cols_diff(R, Rows).

cols_diff([], _).
cols_diff([_|Cs], Rows) :-
	firsts_rest(Rows, Firsts, Rests),
	all_diff(Firsts),
	cols_diff(Cs, Rests).

firsts_rest([], [], []).
firsts_rest([R|Rs], [F|Fs], [Rest|Rests]) :-
	R = [F|Rest],
	firsts_rest(Rs, Fs, Rests).