/* 

  Col sum puzzle in Picat.

  (Since I don't know the name of this puzzle, I call it col sum puzzle.)

  Given the column, row and diagonal sum and some hints, find out
  the rest of the numbers in the matrix.

  One of the sources of this problem is:
  http://mathgrant.blogspot.com/2010/10/grants-review-corner-volume-2.html

  This program was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Picat page: http://www.hakank.org/picat/

*/

% import sat.
import cp.

main => go.

go =>
  nolog,
  % data(1,N,MaxVal,RowSums,ColSums,DiagSums,Clues,AllDiff),
  % data(2,N,MaxVal,RowSums,ColSums,DiagSums,Clues,AllDiff),
  data(3,N,MaxVal,RowSums,ColSums,DiagSums,Clues,AllDiff),
 
  col_sum_puzzle(N,MaxVal,RowSums,ColSums,DiagSums,Clues,AllDiff, X),
  println("X:"),
  foreach(I in 1..N)
    foreach(J in 1..N)
      printf("%2d ", X[I,J])
    end,
    nl
  end,
  nl,
  fail,

  nl.

col_sum_puzzle(N,MaxVal,RowSums,ColSums,DiagSums,Clues,AllDiff, X) => 

  X = new_array(N, N),
  X :: 1..MaxVal, 

  if AllDiff then
    all_different(X.vars)
  end,

  % Clues
  foreach(I in 1..N, J in 1..N,Clues[I,J] > 0)
    X[I,J] #= Clues[I,J]
  end,

  % rowsum/colsum
  foreach(I in 1..N)
    sum([X[I,J] : J in 1..N]) #= RowSums[I],
    sum([X[J,I] : J in 1..N]) #= ColSums[I]
  end,
  
  % diagonal 1
  sum([X[I,I] : I in 1..N]) #= DiagSums[1],
  if DiagSums.len == 2 then
    % diagonal 2
    sum([X[I,N-I+1] : I in 1..N]) #= DiagSums[2]
  end,

  solve($[ff,split],X).

% I'm not sure about the source of this problem,
% It's not exactly the same as the two other instances,
% e.g. it does not enforce all_different, and has
% the valid values 1..9 (not 1..16).
% Also, it has a diagonal 2 hint.
data(1,N,MaxVal,RowSums,ColSums,DiagSums,Clues,AllDiff) =>
  N = 4,
  MaxVal = 9,
  RowSums = [22, 26, 31, 25],
  ColSums = [24, 18, 31, 31],
  DiagSums = [24, 24],

  % the clues (> 0)
  Clues = [[5, 0, 0, 0],
           [0, 0, 8, 0],
           [0, 0, 0, 8],
           [0, 7, 0, 0]],
  AllDiff = false.


% http://mathgrant.blogspot.com/2010/10/grants-review-corner-volume-2.html
% A unique solution:
% [3,2,6]
% [8,4,1]
% [5,9,7]
% 
data(2,N,MaxVal,RowSums,ColSums,DiagSums,Clues,AllDiff) =>
  N = 3,
  MaxVal = N*N,
  RowSums = [11, 13, 21],
  ColSums = [16, 15, 14],
  DiagSums = [14],

  % the clues (> 0)
  Clues = [[0, 2, 0],
           [0, 0, 1],
           [5, 0, 0]],
  AllDiff = true.

% http://mathgrant.blogspot.com/2010/10/grants-review-corner-volume-2.html
% Many solutions (and that is one major complaint that MathGrant writes about).
data(3,N,MaxVal,RowSums,ColSums,DiagSums,Clues,AllDiff) =>
  N = 6,
  MaxVal = N*N,
  RowSums = [60,146,97,117,91,155],
  ColSums = [76,113,137,98,120,122],
  DiagSums = [158],
  Clues = [[ 0, 0, 0, 4, 0, 0],
           [ 0, 0, 0, 0,15, 0],
           [ 0,28, 0, 0, 0, 0],
           [ 0, 0,18, 0, 0, 0],
           [ 6, 0, 0, 0, 0, 0],
           [ 0, 0, 0, 0, 0,36]],
  AllDiff = true.