:-use_module(library(clpfd)).
:-use_module(library(lists)).
:-use_module(library(terms)).
%%% Find a binary code (N,D) with K elements.
code_gen(N,D,K) :-
    M is N*K, length(Code,M),
    domain(Code,0,1),  %%%%%%%%%%% OPTIMIZATIONS %%%%%%%%%%%%%%%%%%%%
    first_word(N,Code),  %%% [0,...,0] is a word
    lexicographic(Code,N),     %%% imposes the lexicographic ordering
    columns_constraints(Code,N,K), %%% imposes two constraints on columns
    constraints(Code,D,N),
    write(N,N,Code),nl,
    term_variables(Code,List),length(List,L),
    write(L),write(' variables.'),nl,
    statistics(runtime,_),
    labeling([ff],Code),!,
    statistics(runtime,[_,T]),
    write('Solution found (msecs): '),write(T),nl,
    write(N,N,Code),nl.
code_gen(_,_,_) :-
    statistics(runtime,[_,T]),
    write('No solution (msecs): '),
    write(T),nl,fail.
%%% Constraints
constraints(Code,_,N) :-
    length(Code,CL),
    CL =< N, !.
constraints(Code,D,N) :-
    extract(N,Code,W1,RestCode),
    allpairs(RestCode,W1,D,N),
    constraints(RestCode,D,N).
allpairs([],_,_,_).
allpairs(Code,Current,D,N) :-
   extract(N,Code,First,Last),
   pair(Current,First,D),
   allpairs(Last,Current,D,N).
pair(Word1,Word2,D) :-
   diffvars(Word1,Word2,Diff),
   domain(Diff,0,1),
   sum(Diff,'#>=',D).
diffvars([],[],[]).
diffvars([A|R],[B|S],[D|T]) :-
   D #= abs(A-B),
   diffvars(R,S,T).
%%%%%% OPTIMIZATION
first_word(N,Code) :-
    extract(N,Code,W1,_),
    set_bits(W1,0).
second_word(N,Code,D) :-
    extract(N,Code,_,RestCode),
    extract(N,RestCode,W2,_),
    length(U,D),
    append(Z,U,W2),
    set_bits(Z,0),
    set_bits(U,1).

lexicographic(Code,N) :-
     length(Code,M),
     M < 2*N, !.
lexicographic(Code,N) :-
     extract(N,Code,W1,RestCode),
     extract(N,RestCode,W2,_),
     %%% Two options. The standard definition (with an "or") and
     %%% a numeric one. Running times are comparable in SICStus.
     lexico(W1,W2,1), %%% first option.
     %bin_value(W1,C1), bin_value(W2,C2), C1 #< C2, %%% second option.
     lexicographic(RestCode,N).

lexico([],[],1).
lexico([A|R],[B|S],C) :-
    lexico(R,S,Val),
    (A #< B #\/ (A #= B #/\ Val)) #<=> C.
columns_constraints(Code,N,K) :-
%    column_constraints(Code,1,N), %%% OPTIONAL
    first_column(Code,N,K).
column_constraints(_,N,N):- !.
column_constraints(Code,M,N) :-
    columns(Code,M,N,Col1,Col2),
    sum(Col1,#=,Uni1),
    sum(Col2,#=,Uni2),
    Uni1 #=< Uni2,
    M1 is M + 1,
    column_constraints(Code,M1,N).
columns([],_,_,[],[]).
columns(Code,M,N,[A|R],[B|S]) :-
    M1 is M+1,
    element(M,Code,A),
    element(M1,Code,B),
    extract(N,Code,_,RestCode),
    columns(RestCode,M,N,R,S).

first_column(Code,N,K) :-
    columns(Code,1,N,C1,_),
    K1 is K // 2 + K mod 2,
    extract(K1,C1,First,Last),
    set_bits(First,0),
    set_bits(Last,1). 
%%%%%%% AUXILIARY
bin_value(W,C) :-
   reverse(W,W1),
   value(W1,C).
value([B],B) :- !.
value([B|R],Val) :-
   value(R,V), Val #= 2*V+B.
extract(N,Code,First,Last) :-
   length(First,N),
   append(First,Last,Code).
set_bits([],_).
set_bits([N|R],N) :-
   set_bits(R,N).
%%%%%%% PRINT RESULT:
write(_,_,[]) :- !.
write(0,N,Code) :-
    !,nl, write(N,N,Code).
write(M,N,[A|Code]) :-
    M1 is M -1,
    write(A),write(' '),
    write(M1,N,Code).
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
