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 decompose_observation_table.py


###################################
# Wat dingetjes over Mealy machines


# Voorbeeld: 2n states, input-alfabet 'a' en 'b', outputs [0...n-1]
def rick_koenders_machine(N):
  transition_fun = {((n, False), 'a'): ((n + 1) % N, False) for n in range(N)}
  transition_fun |= {(((n + 1) % N, True), 'a'): (n % N, True) for n in range(N)}
  transition_fun |= {((n, b), 'b'): (n, not b) for b in [False, True] for n in range(N)}
  output_fun = {((n, b), 'a'): n for b in [False, True] for n in range(N)}
  output_fun |= {((n, b), 'b'): 0 for b in [False, True] for n in range(N)}
  initial_state = (0, False)
  inputs = ['a', 'b']
  outputs = [n for n in range(N)]
  return {'transition_fun': transition_fun, 'output_fun': output_fun, 'initial_state': initial_state, 'inputs': inputs, 'outputs': outputs}


def mealy_sem_q(machine, word, state):
  if len(word) == 0:
    return None
  if len(word) == 1:
    return machine['output_fun'][(state, word[0])]
  else:
    return mealy_sem_q(machine, word[1:], machine['transition_fun'][(state, word[0])])


def mealy_sem(machine, word):
  return mealy_sem_q(machine, word, machine['initial_state'])


def print_table(cell, rs, cs):
  first_col_size = max([len(str(r)) for r in rs])
  col_size = 1 + max([len(str(c)) for c in cs])

  print(''.rjust(first_col_size), end='')
  for c in cs:
    print(str(c).rjust(col_size), end='')
  print('')

  for r in rs:
    print(str(r).rjust(first_col_size), end='')
    for c in cs:
      print(cell(r, c).rjust(col_size), end='')
    print('')


################
# Voorbeeld data
machine = rick_koenders_machine(4)  # 8 states

# L* table: Niet noodzakelijk volledig, werkt ook met minder data, maar ik
# weet niet wat voor garanties we dan kunnen geven. Het is sowieso maar de
# vraag of de kolommen voldoende zijn als we projecteren.
rows = ['', 'a', 'aa', 'aaa', 'aaaa', 'b', 'ab', 'aab', 'aaab']
cols = ['a', 'aa', 'aaa', 'aaaa', 'b', 'ab', 'ba', 'abab']

print_table(lambda r, c: str(mealy_sem(machine, r + c)), rows, cols)

# We zoeken 2 componenten met gezamelijke grootte 6 (minder dan 8)
# als de de total_size te laag is => UNSAT => duurt lang
c = 2
total_size = 6


########################
# Encodering naar logica
print('Start encoding')
os = machine['outputs']  # outputs
rids = [i for i in range(c)]  # components
vpool = IDPool()
cnf = CNF()


# Een hulp variabele voor False en True, maakt de andere variabelen eenvoudiger
def var_const(b):
  return vpool.id(('const', 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(('rel', rid, so1, so2))


# De relatie op outputs geeft een relaties op rijen. Deze zijn geindexeerd
# met de woorden uit 'rows'.
def var_row_rel(rid, r1, r2):
  if r1 == r2:
    return var_const(True)

  [sr1, sr2] = sorted([r1, r2])
  return vpool.id(('row_rel', rid, sr1, sr2))


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


# 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 (o)')
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)])

print('- transitivity (r)')
for rid in rids:
  for rx in rows:
    for ry in rows:
      for rz in rows:
        # als rx R ry en ry R rz dan rx R rz
        cnf.append([-var_row_rel(rid, rx, ry), -var_row_rel(rid, ry, rz), var_row_rel(rid, rx, rz)])

# Constraint zodat de relaties samen alle elementen kunnen onderscheiden.
# (Aka: the bijbehorende quotienten zijn joint-injective.)
print('- injectivity (o)')
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])

# Als outputs equivalent zijn, dan ook sommige rijen, en andersom.
print('rel <=> row_rel')
for rid in rids:
  for rx in rows:
    for ry in rows:
      osx = [mealy_sem(machine, rx + c) for c in cols]
      osy = [mealy_sem(machine, ry + c) for c in cols]
      oss = list(zip(osx, osy))

      # (ox1 ~ oy1 and ox2 ~ oy2 and ...) => rx ~ ry
      cnf.append([-var_rel(rid, ox, oy) for (ox, oy) in oss] + [var_row_rel(rid, rx, ry)])

      # rx ~ ry => oxi ~ oyi
      for ox, oy in oss:
        cnf.append([-var_row_rel(rid, rx, ry), var_rel(rid, ox, oy)])

# De constraints die zorgen dat representanten ook echt representanten zijn.
print('- representatives (r)')
for rid in rids:
  for ix, rx in enumerate(rows):
    # Belangrijkste: een element is een representant, of equivalent met een
    # eerder element. We forceren hiermee dat de solver representanten moet
    # kiezen (voor aan de lijst).
    cnf.append([var_row_rep(rid, rx)] + [var_row_rel(rid, rx, ry) for ry in rows[:ix]])
    for ry in rows[:ix]:
      # rx en ry kunnen niet beide een representant zijn, tenzij ze
      # niet gerelateerd zijn.
      cnf.append([-var_row_rep(rid, rx), -var_row_rep(rid, ry), -var_row_rel(rid, rx, ry)])

# 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_row_rep(rid, rx) for rid in rids for rx in rows], total_size, vpool=vpool)
cnf.extend(cnf_optim)


def print_eqrel(rel, xs):
  print_table(lambda r, c: 'Y' if rel(r, c) else '·', xs, xs)


##################################
# 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('sat :-)')

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

  for rid in rids:
    print(f'Relation {rid}:')
    print_eqrel(lambda x, y: model[var_rel(rid, x, y)], os)

  for rid in rids:
    print(f'Row relation {rid}:')
    print_eqrel(lambda x, y: model[var_row_rel(rid, x, y)], rows)

  # print equivalence classes
  count = 0
  for rid in rids:
    print(f'component {rid}')
    # Eerst verzamelen we de representanten
    for r in rows:
      if model[var_row_rep(rid, r)]:
        print(f'- representative row {r}')
        count += 1

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