% battleships.mzn.model
% Ralph Becket <rafe@csse.unimelb.edu.au>
% Wed Dec 10 13:51:12 EST 2008
% vim: ft=zinc ts=4 sw=4 et tw=0
%
% A battleships puzzle involves working out where each of a fleet of ships
% lies on a square grid.  Clues are given in two ways: some parts of the 
% grid may be revealed as containing open water or part of a ship; and the
% total number of parts of ships appearing in a row or column may also be
% specified.  Individual ships are always surrounded by open water, even
% at the corners.
%
% The following example is taken from
% http://wpc.puzzles.com/history/tests/1999qtest/test.htm
%
%   w _ _ _ _ w _ _ _  2    Fleet:
%   S w _ _ _ _ _ _ _  2    1 battleship    SSSS
%   _ _ _ _ _ _ _ _ _  3    2 cruisers      SSS SSS
%   w _ _ _ _ _ _ w _  2    3 destroyers    SS SS SS
%   _ _ _ _ _ _ _ _ _  2    4 submarines    S S S S
%   _ w S w _ _ _ _ _  1
%   _ _ w _ _ _ _ S _  5
%   _ _ _ _ _ _ _ _ _  1
%   _ _ _ _ _ _ _ _ _  2
%
%   4 0 3 2 2 2 1 4 2
%
%   In this model, each square on the grid is a number from 0..4
%   (0 being water, 4 denoting the bows of a battleship).  A Cruiser,
%   for instance, is described by a contiguous pattern 0 1 2 3 0 on
%   the board.  Consider the following neighbourhood:
%
%       a b c
%       d x e
%       f g h
%
%   Precisely one of the following must hold:
%   - x is 0
%   - x is b + 1 and d = e = 0
%   - x is d + 1 and b = g = 0
%
%   We also need to specify constraints on diagonally adjacent squares:
%   - if x < b then d = e = 0
%   - if x < d then b = g = 0
%
%   The column (row) constraints specify how many non-zero elements a
%   column (row) must have.
%
%   The fleet constraints specify how many instances of each positive number
%   must appear on the board.

include "count.mzn";

% Parameter: the length of side of the grid.
%
int: n;

% Parameter: the number of ships of each class.
%
int: n_classes;
set of int: class = 1..n_classes;
array [class] of int: class_sizes;

set of int: sq = {0} union class;

% The row and column sums.
%
array [row] of var 0..n: row_sums;
array [col] of var 0..n: col_sums;

% We extend the board by one in each direction to add a sea border.
%
set of int: row = 1..n;
set of int: col = 1..n;
set of int: ROW = 0..(n + 1);
set of int: COL = 0..(n + 1);
array [ROW, COL] of var sq: a;

% Add the sea border to the board.
%
constraint forall (r in {0, n + 1}, c in COL) (a[r, c] = 0);
constraint forall (r in ROW, c in {0, n + 1}) (a[r, c] = 0);

% Add the ship constraints.
%
constraint
    forall (r in row, c in col) (
        (
            a[r, c] = 0
        )
    \/
        (
            a[r, c] = a[r, c - 1] + 1
        /\  a[r - 1, c] = 0
        /\  a[r + 1, c] = 0
        )
    \/
        (
            a[r, c] = a[r - 1, c] + 1
        /\  a[r, c - 1] = 0
        /\  a[r, c + 1] = 0
        )
    );

% Add the diagonal constraints.
%
constraint
    forall (r in row, c in col) (
        (
            a[r, c] < a[r, c - 1]
        )
    ->
        (
            a[r - 1, c] = 0
        /\  a[r + 1, c] = 0
        )
    );

constraint
    forall (r in row, c in col) (
        (
            a[r, c] < a[r - 1, c]
        )
    ->
        (
            a[r, c - 1] = 0
        /\  a[r, c + 1] = 0
        )
    );

% Add the row and column constraints.
%
constraint
    forall (r in row) (count([a[r, c] | c in col], 0, n - row_sums[r]));

constraint
    forall (c in col) (count([a[r, c] | r in row], 0, n - col_sums[c]));

% Add the fleet constraints (a_flat is the 1D flattenning of the board, a).
%
array [1..(n * n)] of var sq: a_flat = [a[r, c] | r in row, c in col];

constraint
    forall (s in class) (
        count(a_flat, s, sum (i in class where i >= s) (class_sizes[i]))
    );

solve
    :: int_search(a_flat, first_fail, indomain_max, complete)
    satisfy;

output  [   show(a[r, c]) ++ if c = n then "\n" else " " endif
        |   r in row, c in col
        ];