%% Purpose: Eight queens problem solver using HiPE arrays
%% Author:  Kostis Sagonas <kostis@it.uu.se>
%% Date:    4-Jan-2007
%% Compile: erlc +native nqueens_ha.erl
%% Run:     erl -noshell -run nqueens_ha main 8 1000 -s erlang halt
-module(nqueens_ha).
-export([main/1, queens/3]).

loop(N, N, F, A) -> F(N, A);
loop(I, N, F, A) -> A1 = F(I, A), loop(I+1, N, F, A1).

loop_nl(I, N, F) -> loop(I, N, F, ok), io:nl().

%% Output solution to console if Display is true. 
%% Throws 'abort' when counter signals that MaxCount is reached.
output(_Board, _SolCount, false) -> ok;
output( Board,  SolCount, _Display) ->
    io:format("~-8w/", [SolCount]),
    N = hipe_bifs:array_length(Board) - 1,
    loop_nl(0, N, fun(I,_) -> io:format("~2w"  , [I]) end),
    loop_nl(0, N, fun(I,_) -> io:format("~8w: ", [I]), 
                              QI = hipe_bifs:array_sub(Board, I),
                              loop_nl(0, N, fun(J,_) when J =:= QI -> io:put_chars("Q "); 
                                               (_,_)               -> io:put_chars("* ") 
                                            end) 
                  end).
    
% Is the queen at row I not threatened
safe(_Brd, N, I, N) -> I =:= N;                                        % last column
safe( Brd, N, I, J) ->
    case hipe_bifs:array_sub(Brd, I) - hipe_bifs:array_sub(Brd, J) of
	0 -> I =:= J;             % vert/horiz threat
	D -> case abs(D) =:= I - J of
		 true -> I =:= J;  % diag threat
		 false -> safe(Brd, N, I, J+1)
	     end
    end.

% Main algorithm
queens(N, Times, MaxSol) 
  when is_integer(N), is_integer(Times), is_integer(MaxSol), N > 0, Times > 0 ->
    Board = hipe_bifs:array(N, 0),
    F = fun(_,_) -> queens1(Board, MaxSol, 0, Times =:= 1, 0) end,
    TotalSol = (catch loop(0, Times-1, F, 0)),
    io:format("Total: ~w solutions~n", [TotalSol]).
        
queens1(Brd, MaxSol, SolCount, Display, I) ->
    F = fun(J, Sum) -> hipe_bifs:array_update(Brd, I, J), queens2(Brd, MaxSol, Sum, Display, I) end, 
    loop(0, hipe_bifs:array_length(Brd) - 1, F, SolCount).

queens2(Board, MaxSol, SolCount, Display, I) -> 
    N = hipe_bifs:array_length(Board) - 1,
    case safe(Board, N, I, 0) of
	false             -> SolCount;
	true when I =:= N -> SC = SolCount+1,
			     output(Board, SC, Display), 
			     SC =:= MaxSol andalso throw(MaxSol),
			     SC;
	_                 -> queens1(Board, MaxSol, SolCount, Display, I+1)
    end.

% Main function
main([Size, Times]) -> main([Size, Times, "0"]);
main([Size, Times, MaxSol]) when is_list(Size), is_list(Times), is_list(MaxSol) ->
    queens(list_to_integer(Size), list_to_integer(Times), list_to_integer(MaxSol)).
