/* 

  Set partition problem in Picat.

  Problem formulation from
    http://www.koalog.com/resources/samples/PartitionProblem.java.html
  """
   This is a partition problem.
   Given the set S = {1, 2, ..., n}, 
   it consists in finding two sets A and B such that:

    *  A U B = S
    *  |A| = |B|
    * sum(A) = sum(B)
    * sum_squares(A) = sum_squares(B)
  """


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

*/

import cp.

main => time2(go).


go =>
  N = 20,
  NumSets = 2,
  N2 = N div NumSets, % number of integers in each set

  if N mod 4 != 0 then
     println("N must be a multiple of 4"),
     halt
  end,

  % which set (1..NumSets) is 1..N assigned to?
  A = new_list(N),
  A :: 1..NumSets,

  % number of integers in set S
  Counts = new_list(NumSets), 
  Counts :: N2..N2,

  Sums = new_list(NumSets),
  Sums :: 0..N*N,
  SumSquared = new_list(NumSets),
  SumSquared :: 0..N*N*N,

  foreach(S in 1..NumSets)
    count(S,A,#=,Counts[S]),
    Sums[S] #= sum([I*(A[I]#=S) : I in 1..N]),
    SumSquared[S] #= sum([I*I*(A[I]#=S) : I in 1..N])
  end,
  
  foreach(S in 1..NumSets-1)
    Sums[S] #= Sums[S+1],
    SumSquared[S] #= SumSquared[S+1]
  end,

  % symmetry breaking
  A[1] #= 1,

  % solve($[rand_var,updown],A ++ Sums ++ SumSquared ++ Counts), % might be fast on a single run
  % solve($[degree,updown],A ++ Sums ++ SumSquared ++ Counts),
  solve($[degree,inout],A ++ Sums ++ SumSquared ++ Counts),


  foreach(S in 1..NumSets) 
    printf("set %d: %w\n", S, [I:I in 1..N, A[I] = S])
  end,
  println(sums=Sums),
  println(sumSquared=SumSquared),
  println(counts=Counts),
  
  nl.