from pysat.solvers import Solver
from pysat.card import CardEnc
from pysat.formula import IDPool
from pysat.formula import CNF

### Gebruik:
# Stap 1: pip3 install python-sat
# Stap 2: python3 decomp-sat.py

# TODO: Een L* tabel introduceren, en het aantal representanten van de rijen
# minimaliseren, ipv representanten an sich.

# Voorbeeld data
# snel voorbeeld: n =  27, c = 3 en total_size =  9
# langzaam vb.:   n = 151, c = 4 en total_size = 15
n = 151
c = 4
total_size = 15 # als deze te laag is => UNSAT => duurt lang

os = [i for i in range(n)] # outputs
rids = [i for i in range(c)] # components

# Optimale decompositie met slechts verzamelingen. Zie ook A000792 in de OEIS
# (voor de omkering size -> n), dit groeit met de derdemacht. Dus size is
# ongeveerde de derdemachts-wortel van n. Ik weet niet hoeveel c moet zijn.
# Ook relevant: A007600, Katona's problem en Edmonds' problem.
# n    =  1 2 3 4 5 6 7-9 10-12 13-18 19-27 28-36 37-54 55-81 82-108 109-162
#     ------------------------------------------------------------------------
# c    =* 1 1 1 1 1 2  2    2     3     3     3     4     4     4       5
# size =  1 2 3 4 5 5  6    7     8     9    10    11    12    13      14
#
# *) de c is voor de grootste in het interval


print('Start encoding')
vpool = IDPool()
cnf = CNF()

# Een hulp variabele voor False en True, maakt de andere variabelen eenvoudiger
def var_const(b):
  return(vpool.id(('bool', b)))

cnf.append([var_const(True)])
cnf.append([-var_const(False)])

# Voor elke relatie en elke twee elementen o1 en o2, is er een variabele die
# aangeeft of o1 en o2 gerelateerd zijn. Er is 1 variabele voor xRy en yRx, dus
# symmetrie is al ingebouwd. Reflexiviteit is ook ingebouwd.
def var_rel(rid, o1, o2):
  if o1 == o2:
    return var_const(True)

  [so1, so2] = sorted([o1, o2])
  return(vpool.id(('var_rel', rid, so1, so2)))

# Voor elke relatie, en elke equivalentie-klasse, kiezen we precies 1 element
# als representant. Deze variabele geeft aan welk element.
def var_rep(rid, o):
  return(vpool.id(('var_rep', rid, o)))

# Contraints zodat de relatie een equivalentie relatie is. We hoeven alleen
# maar transitiviteit te encoderen, want refl en symm zijn ingebouwd in de var.
# Dit stukje van het encoderen duurt het langst.
print('- transitivity')
for rid in rids:
  for xo in os:
    for yo in os:
      for zo in os:
        # als xo R yo en yo R zo dan xo R zo
        cnf.append([-var_rel(rid, xo, yo), -var_rel(rid, yo, zo), var_rel(rid, xo, zo)])

# Constraint zodat de relaties samen alle elementen kunnen onderscheiden.
# (Aka: the bijbehorende quotienten zijn joint-injective.)
print('- injectivity')
for xi, xo in enumerate(os):
  for yo in os[xi+1:]:
    # Tenminste een rid moet een verschil maken
    cnf.append([-var_rel(rid, xo, yo) for rid in rids])

# De constraints die zorgen dat representanten ook echt representanten zijn.
print('- representatives')
for rid in rids:
  for xi, xo in enumerate(os):
    # Belangrijkste: een element is een representant, of equivalent met een
    # later element. We forceren hiermee dat de solver representanten moet
    # kiezen (achter aan de lijst).
    cnf.append([var_rep(rid, xo)] + [var_rel(rid, xo, yo) for yo in os[xi+1:]] )
    for _, yo in enumerate(os[xi+1:], xi+1):
      # xo en yo kunnen niet beide een representant zijn, tenzij ze
      # niet gerelateerd zijn.
      cnf.append([-var_rep(rid, xo), -var_rep(rid, yo), -var_rel(rid, xo, yo)])

# Tot slot willen we weinig representanten. Dit doen we met een "atmost"
# formule. Idealiter zoeken we naar de total_size, maar die staat nu vast.
print('- representatives at most k')
cnf_optim = CardEnc.atmost([var_rep(rid, xo) for rid in rids for xo in os], total_size, vpool=vpool)
cnf.extend(cnf_optim)

# Probleem oplossen met solver :-).
print('Start solving')
print('- copying formula')
with Solver(bootstrap_with=cnf) as solver:
  print('- actual solve')
  sat = solver.solve()

  if not sat:
    print('unsat :-(')
    exit()

  print(f'sat :-)')

  # Even omzetten in een makkelijkere data structuur
  print('- get model')
  m = solver.get_model()
  model = {}
  for l in m:
    if l < 0: model[-l] = False
    else: model[l] = True

  # print equivalence classes
  count = 0
  for rid in rids:
    # Eerst verzamelen we de representanten
    reps = []
    for xo in os:
      if model[var_rep(rid, xo)]:
        reps.append(xo)
        count += 1

    # Dan zoeken we wat bij de representant hoort en printen dat.
    for r in reps:
      rep_class = [r]
      for yo in os:
        if yo != r and model[var_rel(rid, r, yo)]:
          rep_class.append(yo)
      print(f'{rid} => class of {r} = {sorted(rep_class)}')

  # count moet gelijk zijn aan cost
  print(f'total size = {count}')