In [52]:
import Numberjack

import forge
from data.logic import _util


Beth = 0
Charles = 1
David = 2
Frank = 3
Jessica = 4
Karen = 5
Taylor = 6

model = Numberjack.Model()

names = ['Beth', 'Charles', 'David', 'Frank', 'Jessica', 'Karen', 'Taylor'] 
people = list(range(7))
grid = []
for x in range(len(names)):
  row = []
  grid.append(row)
  for y in range(len(names)):
    row.append(Numberjack.Variable(0, 6, '%s_%s' % (x, y)))
  model.add(Numberjack.AllDiff(row))
  model.add(row[x] == 0)

#1
model.add(grid[Beth][Taylor] == grid[Jessica][Taylor])
model.add(grid[Jessica][Taylor] == grid[Karen][Taylor])
model.add(grid[Karen][Taylor] == grid[Frank][Beth])

#2
model.add(grid[Charles][Jessica] == grid[David][Jessica])
model.add(grid[David][Jessica] == grid[Frank][Jessica])
model.add(Numberjack.AllDiff([
  grid[Charles][Taylor],
  grid[David][Taylor],
  grid[Frank][Taylor],
]))

#3
model.add(grid[Beth][Charles] == grid[Charles][Beth])
model.add(grid[Beth][David] == grid[David][Beth])
model.add(grid[Charles][David] == grid[David][Charles])

#4
model.add(grid[Beth][Jessica] == grid[David][Taylor])
model.add(grid[Beth][Karen] == grid[Charles][Taylor])

#5
model.add(grid[Frank][Charles] == grid[Beth][Charles])

#6a
model.add(Numberjack.Sum(grid[Beth][Frank] == grid[x][Frank] for x in people) == 3)
#6b
model.add(grid[Beth][Frank] == grid[Taylor][Beth])

#7
model.add(grid[Frank][Karen] == grid[Jessica][Karen])
model.add(grid[Jessica][Karen] == grid[Karen][David])
model.add(grid[Karen][David] == grid[Taylor][David])

#8
model.add(grid[David][Karen] == grid[Taylor][Karen])
model.add(grid[Taylor][Karen] == grid[Jessica][David])
model.add(grid[Jessica][David] == grid[Karen][Beth])

#9
model.add(grid[Karen][Charles] == grid[Taylor][Charles])
model.add(grid[Taylor][Charles] == grid[Jessica][Beth])
model.add(grid[Jessica][Beth] == grid[Frank][David])

#10
model.add(Numberjack.Sum(grid[Karen][Beth] == grid[x][Beth] for x in people) == 2)

#11
model.add(grid[Karen][Jessica] == grid[Taylor][Frank])

#12 ???

#13
def has_n(n):
  sums = []
  for y in people:
    sums.append(Numberjack.Disjunction([grid[x][y] == n for x in people]))
  return Numberjack.Sum(sums)

#13
model.add(has_n(1) > has_n(2))
for i in range(3, 5+1):
  model.add(has_n(2) > has_n(i))
#14
for i in range(3, 5):
  model.add(has_n(i) > has_n(5))
model.add(has_n(5) > has_n(6))

#15 ...

solver = model.load('Mistral')
solver.solve()


if solver.is_unsat():
  print('Impossible')
else:
  for i in range(1, 7):
    print(i, _util.numberjack_solution(has_n(i)))
  for y, row in enumerate(grid):
    result = [names[y]]
    for x, variable in enumerate(row):
      result.append(str(variable.get_value()))
    print('\t'.join(result))


1 7
2 6
3 5
4 5
5 4
6 3
Beth	0	5	1	6	3	2	4
Charles	5	0	4	3	6	1	2
David	1	4	0	2	6	5	3
Frank	4	5	2	0	6	3	1
Jessica	2	1	5	6	0	3	4
Karen	5	2	3	6	1	0	4
Taylor	6	2	3	1	4	5	0

In [ ]: