from pysat.examples.fm import FM
from pysat.formula import IDPool
from pysat.formula import WCNF

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

vpool = IDPool()
wcnf = WCNF()

# Voor elke relatie en elke twee elementen o1 en o2, is er een variabele die
# aangeeft of o1 en o2 gerelateerd zijn. Hierbij moeten o1 en o2 verschillen.
# En er is 1 variabele voor xRy en yRx, dus symmetrie is al ingebouwd.
def var_rel(rid, o1, o2):
  [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)))

def append_hard(clause):
  wcnf.append(clause)

def append_soft(clause):
  wcnf.append(clause, 1)

# Contraints zodat de relatie een equivalentie relatie is. We hoeven alleen
# maar transitiviteit te encoderen.
def append_eq_rel(rid, os):
  for xo in os:
    for yo in os:
      if yo == xo:
        continue
      for zo in os:
        if zo == xo or zo == yo:
          continue
        append_hard([-var_rel(rid, xo, yo), -var_rel(rid, yo, zo), var_rel(rid, xo, zo)])

def append_eq_rel_all(rids, os):
  for rid in rids:
    append_eq_rel(rid, os)

# Constraint zodat de relaties samen alle elementen kunnen onderscheiden.
# (Aka: the bijbehorende quotienten zijn joint-injective.)
def append_injective(rids, os):
  for xi, xo in enumerate(os):
    for _, yo in enumerate(os[xi+1:], xi+1):
      append_hard([-var_rel(rid, xo, yo) for rid in rids])

# De constraints die zorgen dat representanten ook echt representanten zijn.
def append_reps(rids, os):
  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).
      append_hard([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.
        append_hard([-var_rep(rid, xo), -var_rep(rid, yo), -var_rel(rid, xo, yo)])


# Voorbeeld data
os = ['x', 'y', 'z', 'w', 't'] # outputs
rids = [1, 2] # components

# alle constraints toevoegen
append_eq_rel_all(rids, os)
append_injective(rids, os)
append_reps(rids, os)

# We willen zo min mogelijk representanten. Dwz we willen zo veel mogelijk
# elementen die geen representant zijn. Dat doen we met soft clausules en
# MaxSAT optimaliseert dat.
for rid in rids:
  for xo in os:
    append_soft([-var_rep(rid, xo)])


# Probleem oplossen met solver :-). Het PySat pakket bevat nog andere
# MaxSAT solvers: RC2 en LSU. Geen idee welke het beste is. Deze lijkt te
# werken. RC2 kan ook nog meerdere oplossingen enumereren.
with FM(wcnf) as fm:
  if not fm.compute():
    print('unsat :-(')
    exit

  print(f'sat :-) with cost {fm.cost}')

  # Even omzetten in een makkelijkere data structuur
  m = fm.model
  model = {}
  for l in m:
    if l < 0:
      model[-l] = False
    else:
      model[l] = True

  # print equivalence relations
  # for rid in rids:
  #   for xi, xo in enumerate(os):
  #     for _, yo in enumerate(os[xi+1:], xi+1):
  #       if model[var_rel(rid, xo, yo)]:
  #         print(f'{rid} => {xo} == {yo}')
  #       else:
  #         print(f'{rid} => {xo} != {yo}')

  # 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}')


# with RC2(cnf) as rc2:
#   for m in rc2.enumerate():
#     print('model {0} has cost {1}'.format(m, rc2.cost))

# with FM(cnf) as fm:
#   m = fm.compute()
#   print('model {0} has cost {1}'.format(fm.model, fm.cost))