import argparse


class RKMachine:
  # Rick Koenders came up with this example in the case of N=3.
  # FSM will have 2*N states, but can be decomposed into N + 2 states.
  # For N <= 2, the machine is NOT minimal, i.e., there are equivalent states.
  def __init__(self, N):
    self.initial_state = (0, False)
    self.inputs = ['a', 'b']
    self.outputs = [n for n in range(N)]
    self.states = [(n, b) for b in [False, True] for n in range(N)]

    self.transition_map = {((n, 0), 'a'): ((n + 1) % N, 0) for n in range(N)}
    self.transition_map |= {(((n + 1) % N, 1), 'a'): (n % N, 1) for n in range(N)}
    self.transition_map |= {((n, b), 'b'): (n, not b) for b in [False, True] for n in range(N)}

    self.output_map = {((n, b), 'a'): n for b in [False, True] for n in range(N)}
    self.output_map |= {((n, b), 'b'): 0 for b in [False, True] for n in range(N)}

  def state_to_string(self, s):
    (n, b) = s
    return f's{n}Dn' if b else f's{n}Up'

  def transition(self, s, a):
    return self.transition_map[(s, a)]

  def output(self, s, a):
    return self.output_map[(s, a)]


def fsm_to_dot(name, m):
  def transition_string(s, i, o, t):
    return f'{s} -> {t} [label="{i}/{o}"]'

  ret = f'digraph {name} {{\n'
  for s in m.states:
    for a in m.inputs:
      t = m.transition(s, a)
      o = m.output(s, a)
      ret += '  '
      ret += transition_string(m.state_to_string(s), a, o, m.state_to_string(t))
      ret += '\n'

  ret += '}\n'

  return ret


parser = argparse.ArgumentParser()
parser.add_argument('machine', nargs='?', default='rk', choices=['rk'])
parser.add_argument('-n', type=int, default=3, help='size parameter')
args = parser.parse_args()

if args.machine == 'rk':
  print(fsm_to_dot(f'RKMachine_{args.n}', RKMachine(args.n)))