% Board is a list of snake segments
%   segment(Color,X-Y,Direction)
% It may also contain empty squares:
%   segment(empty,X-Y,_)

size(Color, 1) :- member(Color, [empty,g]).
size(Color, 2) :- member(Color, [r,y,b]).
size(Color, 3) :- member(Color, [o,i]).
size(Color, 4) :- member(Color, [v]).

next(r, o).
next(o, y).
next(y, g).
next(g, b).
next(b, i).
next(i, v).
next(v, end).

position(X-Y) :- between(1, 5, X), between(1, 4, Y).

direction(left).
direction(right).
direction(up).
direction(down).

solve(Board, Solution) :- solve(r, Board, Solution).

solve(end, Board, Board) :- final_check(Board).
solve(Color, Board, Solution) :-
    next(Color, Color1),
    ( member(segment(Color,Pos,Dir), Board) ->
      connected(segment(Color,Pos,Dir), Board),
      solve(Color1, Board, Solution)
    ; position(Pos),
      direction(Dir),
      Segment = segment(Color,Pos,Dir),
      check(Segment, Board),
      solve(Color1, [Segment|Board], Solution)
    ).

check(Segment, Board) :-
    Segment = segment(Color,Pos,Dir),
    size(Color, Size),
    ( Size =:= 1 -> Dir = left ; true ),
    forall(between(1, Size, K),
           ( step(Pos, Dir, K, Pos1),
             position(Pos1),
             \+ occupied(Pos1, Board, _)
           )),
    connected(Segment, Board).

step(Pos, _, 1, Pos).
step(X-Y, left, N, Pos) :-
    N > 1, N1 is N - 1,
    X1 is X - 1,
    step(X1-Y, left, N1, Pos).
step(X-Y, right, N, Pos) :-
    N > 1, N1 is N - 1,
    X1 is X + 1,
    step(X1-Y, right, N1, Pos).
step(X-Y, up, N, Pos) :-
    N > 1, N1 is N - 1,
    Y1 is Y - 1,
    step(X-Y1, up, N1, Pos).
step(X-Y, down, N, Pos) :-
    N > 1, N1 is N - 1,
    Y1 is Y + 1,
    step(X-Y1, down, N1, Pos).

occupied(Target, Board, Color) :-
    member(segment(Color,Pos,Dir), Board),
    size(Color, Size),
    between(1, Size, K),
    step(Pos, Dir, K, Target).

connected(segment(r,_,_), _).
connected(segment(Color,Pos,_), Board) :-
    next(Prev, Color),
    member(segment(Prev,Pos0,Dir0), Board),
    size(Prev, Size),
    step(Pos0, Dir0, Size, Last),
    adjacent(Last, Pos).

adjacent(X1-Y, X2-Y) :- abs(X1 - X2) =:= 1.
adjacent(X-Y1, X-Y2) :- abs(Y1 - Y2) =:= 1.

final_check(Board) :-
    findall(Pos,
            ( position(Pos),
              \+ ( occupied(Pos, Board, Color), Color \= empty )
            ),
            Empty),
    separate(Empty).

separate([_]).
separate([Pos|Rest]) :-
    \+ ( member(Pos1, Rest), adjacent(Pos, Pos1) ),
    separate(Rest).

show(Board) :- show(1, 1, Board), !.
show(5, _, _).
show(Row, 6, Board) :- Row1 is Row + 1, nl, show(Row1, 1, Board).
show(Row, Col, Board) :-
    Col1 is Col + 1,
    ( occupied(Col-Row, Board, Color), Color \= empty ->
      write(Color), write(' ')
    ; write('. ')
    ),
    show(Row, Col1, Board).

take(0, _, []).
take(N, [X|Xs], [X|Ys]) :- N > 0, N1 is N - 1, take(N1, Xs, Ys).
take(N, [_|Xs], Ys) :- N > 0, take(N, Xs, Ys).

% n non-empty, k empty fixed
find_n_k(N, K, Solution) :-
    solve([], Solution),
    findall(segment(empty,Pos,left),
            ( position(Pos), \+ occupied(Pos, Solution, _) ),
            Empty),
    take(N, Solution, L1),
    take(K, Empty, L2),
    append(L1, L2, Board),
    findall(S, solve(Board, S), [_]),
    write(Board), nl, show(Solution), nl.

% unambiguous puzzle for the given solution with N parts
try_n(N, Solution) :-
    take(N, Solution, Puzzle),
    findall(S, solve(Puzzle, S), [_]),
    write(Puzzle), nl, show(Solution), nl.