"""Data generation script for logic programs."""
import argparse
import random as R

# Symbol Pool
CONST_SYMBOLS = "abcdefghijklmnopqrstuvwxyz"
VAR_SYMBOLS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
PRED_SYMBOLS = "abcdefghijklmnopqrstuvwxyz"
EXTRA_SYMBOLS = "-,()"

CHARS = sorted(list(set(CONST_SYMBOLS+VAR_SYMBOLS+PRED_SYMBOLS+EXTRA_SYMBOLS)))
# Reserve 0 for padding
CHAR_IDX = dict((c, i+1) for i, c in enumerate(CHARS))
IDX_CHAR = [0]
IDX_CHAR.extend(CHARS)

# Predicate Templates
FACT_T = "{}."
RULE_T = "{}:-{}."
PRED_T = "{}({})"
ARG_SEP = ','
PRED_SEP = ';'
TARGET_T = "? {} {} {}"

def choices(symbols, k):
  """Return k many symbols with replacement. Added in v3.6."""
  return [R.choice(symbols) for _ in range(k)]

def r_string(symbols, length):
  """Return random sequence from given symbols."""
  return ''.join(R.choice(symbols)
                 for _ in range(length))

def r_symbols(size, symbols, length, used=None):
  """Return unique random from given symbols."""
  if length == 1 and not used:
    return R.sample(symbols, size)
  rset, used = set(), set(used or [])
  while len(rset) < size:
    s = r_string(symbols, R.randint(1, length))
    if s not in used:
      rset.add(s)
  return list(rset)

def r_consts(size, used=None):
  """Return size many unique constants."""
  return r_symbols(size, CONST_SYMBOLS, ARGS.constant_length, used)

def r_vars(size, used=None):
  """Return size many unique variables."""
  return r_symbols(size, VAR_SYMBOLS, ARGS.variable_length, used)

def r_preds(size, used=None):
  """Return size many unique predicates."""
  return r_symbols(size, PRED_SYMBOLS, ARGS.predicate_length, used)

def write_p(pred):
  """Format single predicate tuple into string."""
  return PRED_T.format(pred[0], ARG_SEP.join(pred[1]))

def write_r(preds):
  """Convert rule predicate tuple into string."""
  head = write_p(preds[0])
  # Is it just a fact
  if len(preds) == 1:
    return FACT_T.format(head)
  # We have a rule
  return RULE_T.format(head, PRED_SEP.join([write_p(p) for p in preds[1:]]))

def output(context, targets):
  """Print the context and given targets."""
  # context: [[('p', ['a', 'b'])], ...]
  # targets: [(('p', ['a', 'b']), 1, [0,1,2]), ...]
  print('\n'.join([write_r(c) for c in context]))
  for t, v, s in targets:
    print(TARGET_T.format(write_r([t]), v, ','.join(map(str, s))))

def gen_task(context, targets, upreds):
  """Fill context with random preds and output program."""
  # Fill with random rules up to certain task
  ctx = context.copy() # Don't modify the original context
  for _ in range(ARGS.noise_size):
    rtask = R.randint(1, max(1, ARGS.task))
    rtask = 1 if ARGS.arity != 2 and rtask == 8 else rtask
    ntask = "gen_task" + str(rtask)
    ctx.append(globals()[ntask](upreds))
  output(ctx, targets)

def add_pred(context, pred, upreds, uconsts, psuccess=0.0):
  """Fail a ground case predicate given context."""
  # Maybe succeed by adding to context
  if R.random() < psuccess:
    context.append([pred])
  if R.random() < 0.5:
    # The constant doesn't match
    args = pred[1].copy()
    args[R.randrange(len(args))] = r_consts(1, uconsts)[0]
    context.append([(pred[0], args)])
  if R.random() < 0.5:
    # The predicate doesn't match
    p = r_preds(1, upreds)[0]
    upreds.append(p)
    context.append([(p, pred[1])])
  # The predicate doesn't appear at all

def gen_task1(upreds=None):
  """Ground instances only: p(a).q(c,b)."""
  # One or two argument predicate
  preds = r_preds(2, upreds)
  args = r_consts(R.randint(1, ARGS.arity))
  rule = [(preds[0], args)]
  if upreds:
    return rule
  ctx = list()
  add_pred(ctx, rule[0], preds, args, 1.0)
  # Successful case when query appears in context
  targets = [(rule[0], 1, [0])]
  # Fail case
  args = r_consts(R.randint(1, 2))
  fpred = (preds[1], args)
  add_pred(ctx, fpred, preds, args)
  targets.append((fpred, 0, [-1]))
  gen_task(ctx, targets, preds)

def gen_task2(upreds=None):
  """Variablised facts only: p(X).q(X,Y)."""
  preds = r_preds(2, upreds)
  ctx, targets = list(), list()
  if R.random() < 0.5 and ARGS.arity > 1:
    # Double variable same argument
    v = r_vars(1)[0]
    rule = [(preds[0], [v, v])]
    if upreds:
      return rule
    ctx.append(rule)
    # Successful double variable grounding
    cs = r_consts(2)
    c = R.choice(cs)
    targets.append(((preds[0], [c, c]), 1, [0]))
    # Fail on non-unique variable grounding
    targets.append(((preds[0], cs), 0, [-1]))
  else:
    # Double variable different argument
    # Single variable argument
    argc = R.randint(1, ARGS.arity)
    args = r_vars(argc)
    rule = [(preds[0], args)]
    if upreds:
      return rule
    ctx.append(rule)
    # Successful unique argument grounding
    args = choices(r_consts(2), argc)
    targets.append(((preds[0], args), 1, [0]))
    # Fail on out of context predicate with same arguments
    targets.append(((preds[1], args), 0, [-1]))
  gen_task(ctx, targets, preds)

def nstep_deduction(steps, negation=False, upreds=None):
  """Generate nstep deduction programs given number of steps."""
  assert steps >= 1, "Need at least 1 step deduction."
  preds = r_preds(2 if upreds else 3+steps, upreds)
  consts = r_consts(2)
  ctx, targets = list(), list()
  prefix = '-' if negation else ''
  if R.random() < 0.5 and ARGS.arity > 1:
    # Double variable swap deduction rules
    vs = r_vars(2)
    rule = [(preds[0], vs), (prefix+preds[1], vs[::-1])]
    if upreds:
      return rule
    ctx.append(rule)
    # Add the n steps
    swapc = 1
    for j in range(steps-1):
      vs = r_vars(2)
      toswap = R.random() < 0.5 # Do we swap again?
      args = vs[::-1] if toswap else vs
      ctx.append([(preds[j+1], vs), (preds[j+2], args)])
      swapc += int(toswap)
    # Add the ground case
    args = r_consts(2)
    add_pred(ctx, (preds[steps], args), preds, consts, 1.0)
    args = args if swapc % 2 == 0 else args[::-1]
    succ_supps = list(range(steps+1))
    fail_supps = succ_supps[:-1] + [-1]
    targets.append(((preds[0], args), 1-int(negation), succ_supps))
    targets.append(((preds[0], args[::-1]), int(negation), fail_supps))
    gen_task(ctx, targets, preds)
  else:
    # Double variable non-swap deduction rules
    # Single variable deduction rules
    argc = R.randint(1, ARGS.arity)
    vs = r_vars(argc)
    rule = [(preds[0], vs), (prefix+preds[1], vs)]
    if upreds:
      return rule
    ctx.append(rule)
    # Add the n steps
    for j in range(steps-1):
      vs = r_vars(argc)
      ctx.append([(preds[j+1], vs), (preds[j+2], vs)])
    args = choices(r_consts(2), argc)
    # Add the ground case
    cctx = ctx.copy()
    spred = (preds[steps], args)
    add_pred(cctx, spred, preds, args, 1.0)
    succ_supps = list(range(steps+1))
    targets = [((preds[0], args), 1-int(negation), succ_supps)]
    gen_task(cctx, targets, preds)
    # Add failure case
    if R.random() < 0.5:
      # Fail on broken chain
      p = r_preds(1, preds)[0]
      preds.append(p)
      add_pred(ctx, spred, preds, args, 1.0)
      ctx[0] = [(preds[0], vs), (prefix+p, vs)]
      supps = [0] + [-1]*steps
    else:
      # Fail on last ground case
      add_pred(ctx, spred, preds, args)
      supps = succ_supps[:-1] + [-1]
    targets = [((preds[0], args), int(negation), supps)]
    gen_task(ctx, targets, preds)

def gen_task3(upreds=None):
  """Single step deduction: p(X):-q(X)."""
  return nstep_deduction(1, upreds=upreds)

def gen_task4(upreds=None):
  """Double step deduction: p(X):-q(X).q(X):-r(X)."""
  return nstep_deduction(2, upreds=upreds)

def gen_task5(upreds=None):
  """Triple step deduction."""
  return nstep_deduction(3, upreds=upreds)

def logical_and(negation=False, upreds=None):
  """Logical AND with optional negation: p(X):-q(X);r(X)."""
  preds = r_preds(3, upreds)
  argc = R.randint(1, ARGS.arity)
  # Double variable AND with different vars
  # Single variable AND
  vs = r_vars(argc)
  rule = [(preds[0], vs),
          (preds[1], vs[:1]),
          (preds[2], vs[1:] or vs)]
  if upreds:
    return rule
  ctx = [rule]
  # Create the ground arguments
  args = choices(r_consts(2), argc)
  prem1 = (preds[1], args[:1])
  prem2 = (preds[2], args[1:] or args)
  prems = [prem1, prem2]
  if negation:
    # Add negation to random predicate in body
    ridx = R.randrange(2)
    p, pargs = ctx[-1][ridx+1]
    ctx[-1][ridx+1] = ('-' + p, pargs)
    # Successful case when negation fails
    cctx = ctx.copy()
    add_pred(cctx, prems[ridx], preds, args)
    supps = [0, -1, -1] if ridx == 0 else [0, len(cctx), -1]
    cctx.append([prems[1-ridx]])
    targets = [((preds[0], args), 1, supps)]
    gen_task(cctx, targets, preds)
    # Fail one premise randomly
    fidx = R.randrange(2)
    if ridx == fidx:
      # To fail negation add ground instance
      supps = [0, len(ctx), -1]
      ctx.append([prems[ridx]])
      # Succeed other with some probability
      add_pred(ctx, prems[1-ridx], preds, args, 0.8)
    else:
      # Fail non-negated premise
      supps = [0, -1, -1]
      add_pred(ctx, prems[1-ridx], preds, args)
      # Still succeed negation
      add_pred(ctx, prems[ridx], preds, args)
    targets = [((preds[0], args), 0, supps)]
    gen_task(ctx, targets, preds)
  else:
    # Create successful context
    cctx = ctx.copy()
    add_pred(cctx, prems[0], preds, args, 1.0)
    supps = [0, 1, len(cctx)]
    add_pred(cctx, prems[1], preds, args, 1.0)
    targets = [((preds[0], args), 1, supps)]
    gen_task(cctx, targets, preds)
    # Fail one premise randomly
    fidx = R.randrange(2)
    add_pred(ctx, prems[fidx], preds, args)
    # Succeed the other with some probability
    add_pred(ctx, prems[1-fidx], preds, args, 0.8)
    targets = [((preds[0], args), 0, [0, -1, -1])]
    gen_task(ctx, targets, preds)

def gen_task6(upreds=None):
  """Logical AND: p(X):-q(X);r(X)."""
  return logical_and(upreds=upreds)

def logical_or(negation=False, upreds=None):
  """Logical OR with optional negation: p(X):-q(X).p(X):-r(X)."""
  preds = r_preds(3, upreds)
  # Double or single variable OR
  argc = R.randint(1, ARGS.arity)
  vs = r_vars(argc)
  swap = R.random() < 0.5
  prefix = '-' if negation else ''
  rule = [(preds[0], vs), (prefix + preds[1], vs[::-1] if swap else vs)]
  if upreds:
    return rule
  ctx = list()
  ctx.append(rule)
  # Add the extra branching rules
  ctx.append([(preds[0], vs), (preds[2], vs)])
  args = r_consts(argc)
  ctx.append([(preds[0], args)])
  if swap and argc == 2:
    args = r_consts(argc, args)
    add_pred(ctx, (preds[1], args), preds, args, 1.0)
    args = args[::-1] if swap else args
    targets = [((preds[0], args), 1-int(negation), [0, 3]),
               ((preds[0], args[::-1]), int(negation), [0, -1])]
    gen_task(ctx, targets, preds)
  elif not negation and R.random() < 0.2:
    # Sneaky shorcut case
    targets = [((preds[0], args), 1, [2, -1])]
    gen_task(ctx, targets, preds)
    del ctx[-1]
    targets = [((preds[0], args), 0, [-1, -1])]
    gen_task(ctx, targets, preds)
  else:
    # Succeed either from them
    prems = [(preds[i], r_consts(argc, args)) for i in range(1, 3)]
    sidx = R.randrange(2)
    cctx = ctx.copy()
    if negation and sidx == 0:
      # Succeed by failing negation
      supps = [0, -1]
      add_pred(cctx, prems[0], preds, prems[0][1])
      # Possibly succeed other prem
      add_pred(cctx, prems[1], preds, prems[1][1], 0.2)
    else:
      # Succeed by adding ground case
      supps = [sidx, len(cctx)]
      add_pred(cctx, prems[sidx], preds, prems[sidx][1], 1.0)
      # Possibly succeed other prem
      add_pred(cctx, prems[1-sidx], preds, prems[1-sidx][1], 0.2)
    targets = [((preds[0], prems[sidx][1]), 1, supps)]
    gen_task(cctx, targets, preds)
    # Fail both
    add_pred(ctx, prems[0], preds, prems[0][1], int(negation))
    add_pred(ctx, prems[1], preds, prems[1][1])
    targets = [((preds[0], prems[sidx][1]), 0, [-1, -1])]
    gen_task(ctx, targets, preds)


def gen_task7(upreds=None):
  """Logical OR: p(X):-q(X).p(X):-r(X)."""
  return logical_or(upreds=upreds)

def gen_task8(upreds=None):
  """Transitive case: p(X,Y):-q(X,Z);r(Z,Y)."""
  assert ARGS.arity == 2, "Transitive task arity must be 2."
  preds = r_preds(3, upreds)
  # Existential variable with single choice
  vs = r_vars(3)
  rule = [(preds[0], [vs[0], vs[2]]),
          (preds[1], vs[:2]),
          (preds[2], vs[1:])]
  if upreds:
    return rule
  ctx = [rule]
  # Add matching ground cases
  args = r_consts(3)
  add_pred(ctx, (preds[1], args[:2]), preds, args, 1.0)
  succ_supps = [0, 1, len(ctx), -1]
  add_pred(ctx, (preds[2], args[1:]), preds, args, 1.0)
  # Add non-matching ground cases
  argso = r_consts(3)
  argso.insert(R.randint(1, 2), r_consts(1, argso)[0])
  add_pred(ctx, (preds[1], argso[:2]), preds, argso, 0.5)
  add_pred(ctx, (preds[2], argso[2:]), preds, argso, 0.5)
  # Successful case
  # Fail on half-matching existential
  targets = [((preds[0], [args[0], args[2]]), 1, succ_supps),
             ((preds[0], [argso[0], argso[3]]), 0, [0, -1, -1, -1])]
  gen_task(ctx, targets, preds)

def gen_task9(upreds=None):
  """Single step deduction with NBF: p(X):--q(X)."""
  return nstep_deduction(1, True, upreds)

def gen_task10(upreds=None):
  """Double step deduction with NBF: p(X):--q(X).q(X):-r(X)."""
  return nstep_deduction(2, True, upreds)

def gen_task11(upreds=None):
  """Logical AND with NBF: p(X):-q(X);-r(X)."""
  return logical_and(True, upreds)

def gen_task12(upreds=None):
  """Logical OR with NBF: p(X):--q(X).p(X):-r(X)."""
  return logical_or(True, upreds)

def gen_task0():
  """Generate an ILP task example."""
  argc = 1
  goal = 'f'
  premise = 'b'
  ctx, targets = list(), list()
  # Generate according to goal <- premise
  args = r_consts(argc)
  # Add the successful ground case
  ctx.append([(premise, args)])
  targets.append(((goal, args), 1))
  # Fail on non-matching constant
  args = args.copy()
  args[R.randrange(len(args))] = r_consts(1, args)[0]
  preds = r_preds(3)
  ctx.append([(preds[0], args)])
  targets.append(((goal, args), 0))
  # Add padding length dummy rule
  vs = r_vars(argc)
  ctx.append([(preds[1], vs), (preds[2], vs)])
  preds.extend([goal, premise])
  gen_task(ctx, targets, preds)

if __name__ == '__main__':
  # pylint: disable=line-too-long
  # Arguments
  parser = argparse.ArgumentParser(description="Generate logic program data.")
  parser.add_argument("-t", "--task", default=1, type=int, help="The task to generate.")
  parser.add_argument("-s", "--size", default=1, type=int, help="Number of programs to generate.")
  # Configuration parameters
  parser.add_argument("-ns", "--noise_size", default=2, type=int, help="Size of added noise rules.")
  parser.add_argument("-cl", "--constant_length", default=1, type=int, help="Length of constants.")
  parser.add_argument("-vl", "--variable_length", default=1, type=int, help="Length of variables.")
  parser.add_argument("-pl", "--predicate_length", default=1, type=int, help="Length of predicates.")
  parser.add_argument("-ar", "--arity", default=2, type=int, help="Arity.")
  # Task specific options
  parser.add_argument("--nstep", type=int, help="Generate nstep deduction programs.")
  ARGS = parser.parse_args()

  # Generate given task
  task = "gen_task" + str(ARGS.task)
  for _ in range(ARGS.size):
    if ARGS.nstep:
      nstep_deduction(ARGS.nstep)
    else:
      globals()[task]()

12426dadbc64e6950311a700e1fe485546abd303
<mask token>


def choices(symbols, k):
    """Return k many symbols with replacement. Added in v3.6."""
    return [R.choice(symbols) for _ in range(k)]


<mask token>


def r_symbols(size, symbols, length, used=None):
    """Return unique random from given symbols."""
    if length == 1 and not used:
        return R.sample(symbols, size)
    rset, used = set(), set(used or [])
    while len(rset) < size:
        s = r_string(symbols, R.randint(1, length))
        if s not in used:
            rset.add(s)
    return list(rset)


def r_consts(size, used=None):
    """Return size many unique constants."""
    return r_symbols(size, CONST_SYMBOLS, ARGS.constant_length, used)


<mask token>


def r_preds(size, used=None):
    """Return size many unique predicates."""
    return r_symbols(size, PRED_SYMBOLS, ARGS.predicate_length, used)


def write_p(pred):
    """Format single predicate tuple into string."""
    return PRED_T.format(pred[0], ARG_SEP.join(pred[1]))


<mask token>


def output(context, targets):
    """Print the context and given targets."""
    print('\n'.join([write_r(c) for c in context]))
    for t, v, s in targets:
        print(TARGET_T.format(write_r([t]), v, ','.join(map(str, s))))


def gen_task(context, targets, upreds):
    """Fill context with random preds and output program."""
    ctx = context.copy()
    for _ in range(ARGS.noise_size):
        rtask = R.randint(1, max(1, ARGS.task))
        rtask = 1 if ARGS.arity != 2 and rtask == 8 else rtask
        ntask = 'gen_task' + str(rtask)
        ctx.append(globals()[ntask](upreds))
    output(ctx, targets)


def add_pred(context, pred, upreds, uconsts, psuccess=0.0):
    """Fail a ground case predicate given context."""
    if R.random() < psuccess:
        context.append([pred])
    if R.random() < 0.5:
        args = pred[1].copy()
        args[R.randrange(len(args))] = r_consts(1, uconsts)[0]
        context.append([(pred[0], args)])
    if R.random() < 0.5:
        p = r_preds(1, upreds)[0]
        upreds.append(p)
        context.append([(p, pred[1])])


<mask token>


def gen_task2(upreds=None):
    """Variablised facts only: p(X).q(X,Y)."""
    preds = r_preds(2, upreds)
    ctx, targets = list(), list()
    if R.random() < 0.5 and ARGS.arity > 1:
        v = r_vars(1)[0]
        rule = [(preds[0], [v, v])]
        if upreds:
            return rule
        ctx.append(rule)
        cs = r_consts(2)
        c = R.choice(cs)
        targets.append(((preds[0], [c, c]), 1, [0]))
        targets.append(((preds[0], cs), 0, [-1]))
    else:
        argc = R.randint(1, ARGS.arity)
        args = r_vars(argc)
        rule = [(preds[0], args)]
        if upreds:
            return rule
        ctx.append(rule)
        args = choices(r_consts(2), argc)
        targets.append(((preds[0], args), 1, [0]))
        targets.append(((preds[1], args), 0, [-1]))
    gen_task(ctx, targets, preds)


def nstep_deduction(steps, negation=False, upreds=None):
    """Generate nstep deduction programs given number of steps."""
    assert steps >= 1, 'Need at least 1 step deduction.'
    preds = r_preds(2 if upreds else 3 + steps, upreds)
    consts = r_consts(2)
    ctx, targets = list(), list()
    prefix = '-' if negation else ''
    if R.random() < 0.5 and ARGS.arity > 1:
        vs = r_vars(2)
        rule = [(preds[0], vs), (prefix + preds[1], vs[::-1])]
        if upreds:
            return rule
        ctx.append(rule)
        swapc = 1
        for j in range(steps - 1):
            vs = r_vars(2)
            toswap = R.random() < 0.5
            args = vs[::-1] if toswap else vs
            ctx.append([(preds[j + 1], vs), (preds[j + 2], args)])
            swapc += int(toswap)
        args = r_consts(2)
        add_pred(ctx, (preds[steps], args), preds, consts, 1.0)
        args = args if swapc % 2 == 0 else args[::-1]
        succ_supps = list(range(steps + 1))
        fail_supps = succ_supps[:-1] + [-1]
        targets.append(((preds[0], args), 1 - int(negation), succ_supps))
        targets.append(((preds[0], args[::-1]), int(negation), fail_supps))
        gen_task(ctx, targets, preds)
    else:
        argc = R.randint(1, ARGS.arity)
        vs = r_vars(argc)
        rule = [(preds[0], vs), (prefix + preds[1], vs)]
        if upreds:
            return rule
        ctx.append(rule)
        for j in range(steps - 1):
            vs = r_vars(argc)
            ctx.append([(preds[j + 1], vs), (preds[j + 2], vs)])
        args = choices(r_consts(2), argc)
        cctx = ctx.copy()
        spred = preds[steps], args
        add_pred(cctx, spred, preds, args, 1.0)
        succ_supps = list(range(steps + 1))
        targets = [((preds[0], args), 1 - int(negation), succ_supps)]
        gen_task(cctx, targets, preds)
        if R.random() < 0.5:
            p = r_preds(1, preds)[0]
            preds.append(p)
            add_pred(ctx, spred, preds, args, 1.0)
            ctx[0] = [(preds[0], vs), (prefix + p, vs)]
            supps = [0] + [-1] * steps
        else:
            add_pred(ctx, spred, preds, args)
            supps = succ_supps[:-1] + [-1]
        targets = [((preds[0], args), int(negation), supps)]
        gen_task(ctx, targets, preds)


def gen_task3(upreds=None):
    """Single step deduction: p(X):-q(X)."""
    return nstep_deduction(1, upreds=upreds)


def gen_task4(upreds=None):
    """Double step deduction: p(X):-q(X).q(X):-r(X)."""
    return nstep_deduction(2, upreds=upreds)


def gen_task5(upreds=None):
    """Triple step deduction."""
    return nstep_deduction(3, upreds=upreds)


def logical_and(negation=False, upreds=None):
    """Logical AND with optional negation: p(X):-q(X);r(X)."""
    preds = r_preds(3, upreds)
    argc = R.randint(1, ARGS.arity)
    vs = r_vars(argc)
    rule = [(preds[0], vs), (preds[1], vs[:1]), (preds[2], vs[1:] or vs)]
    if upreds:
        return rule
    ctx = [rule]
    args = choices(r_consts(2), argc)
    prem1 = preds[1], args[:1]
    prem2 = preds[2], args[1:] or args
    prems = [prem1, prem2]
    if negation:
        ridx = R.randrange(2)
        p, pargs = ctx[-1][ridx + 1]
        ctx[-1][ridx + 1] = '-' + p, pargs
        cctx = ctx.copy()
        add_pred(cctx, prems[ridx], preds, args)
        supps = [0, -1, -1] if ridx == 0 else [0, len(cctx), -1]
        cctx.append([prems[1 - ridx]])
        targets = [((preds[0], args), 1, supps)]
        gen_task(cctx, targets, preds)
        fidx = R.randrange(2)
        if ridx == fidx:
            supps = [0, len(ctx), -1]
            ctx.append([prems[ridx]])
            add_pred(ctx, prems[1 - ridx], preds, args, 0.8)
        else:
            supps = [0, -1, -1]
            add_pred(ctx, prems[1 - ridx], preds, args)
            add_pred(ctx, prems[ridx], preds, args)
        targets = [((preds[0], args), 0, supps)]
        gen_task(ctx, targets, preds)
    else:
        cctx = ctx.copy()
        add_pred(cctx, prems[0], preds, args, 1.0)
        supps = [0, 1, len(cctx)]
        add_pred(cctx, prems[1], preds, args, 1.0)
        targets = [((preds[0], args), 1, supps)]
        gen_task(cctx, targets, preds)
        fidx = R.randrange(2)
        add_pred(ctx, prems[fidx], preds, args)
        add_pred(ctx, prems[1 - fidx], preds, args, 0.8)
        targets = [((preds[0], args), 0, [0, -1, -1])]
        gen_task(ctx, targets, preds)


<mask token>


def logical_or(negation=False, upreds=None):
    """Logical OR with optional negation: p(X):-q(X).p(X):-r(X)."""
    preds = r_preds(3, upreds)
    argc = R.randint(1, ARGS.arity)
    vs = r_vars(argc)
    swap = R.random() < 0.5
    prefix = '-' if negation else ''
    rule = [(preds[0], vs), (prefix + preds[1], vs[::-1] if swap else vs)]
    if upreds:
        return rule
    ctx = list()
    ctx.append(rule)
    ctx.append([(preds[0], vs), (preds[2], vs)])
    args = r_consts(argc)
    ctx.append([(preds[0], args)])
    if swap and argc == 2:
        args = r_consts(argc, args)
        add_pred(ctx, (preds[1], args), preds, args, 1.0)
        args = args[::-1] if swap else args
        targets = [((preds[0], args), 1 - int(negation), [0, 3]), ((preds[0
            ], args[::-1]), int(negation), [0, -1])]
        gen_task(ctx, targets, preds)
    elif not negation and R.random() < 0.2:
        targets = [((preds[0], args), 1, [2, -1])]
        gen_task(ctx, targets, preds)
        del ctx[-1]
        targets = [((preds[0], args), 0, [-1, -1])]
        gen_task(ctx, targets, preds)
    else:
        prems = [(preds[i], r_consts(argc, args)) for i in range(1, 3)]
        sidx = R.randrange(2)
        cctx = ctx.copy()
        if negation and sidx == 0:
            supps = [0, -1]
            add_pred(cctx, prems[0], preds, prems[0][1])
            add_pred(cctx, prems[1], preds, prems[1][1], 0.2)
        else:
            supps = [sidx, len(cctx)]
            add_pred(cctx, prems[sidx], preds, prems[sidx][1], 1.0)
            add_pred(cctx, prems[1 - sidx], preds, prems[1 - sidx][1], 0.2)
        targets = [((preds[0], prems[sidx][1]), 1, supps)]
        gen_task(cctx, targets, preds)
        add_pred(ctx, prems[0], preds, prems[0][1], int(negation))
        add_pred(ctx, prems[1], preds, prems[1][1])
        targets = [((preds[0], prems[sidx][1]), 0, [-1, -1])]
        gen_task(ctx, targets, preds)


def gen_task7(upreds=None):
    """Logical OR: p(X):-q(X).p(X):-r(X)."""
    return logical_or(upreds=upreds)


<mask token>


def gen_task10(upreds=None):
    """Double step deduction with NBF: p(X):--q(X).q(X):-r(X)."""
    return nstep_deduction(2, True, upreds)


<mask token>


def gen_task12(upreds=None):
    """Logical OR with NBF: p(X):--q(X).p(X):-r(X)."""
    return logical_or(True, upreds)


def gen_task0():
    """Generate an ILP task example."""
    argc = 1
    goal = 'f'
    premise = 'b'
    ctx, targets = list(), list()
    args = r_consts(argc)
    ctx.append([(premise, args)])
    targets.append(((goal, args), 1))
    args = args.copy()
    args[R.randrange(len(args))] = r_consts(1, args)[0]
    preds = r_preds(3)
    ctx.append([(preds[0], args)])
    targets.append(((goal, args), 0))
    vs = r_vars(argc)
    ctx.append([(preds[1], vs), (preds[2], vs)])
    preds.extend([goal, premise])
    gen_task(ctx, targets, preds)


<mask token>

<mask token>


def choices(symbols, k):
    """Return k many symbols with replacement. Added in v3.6."""
    return [R.choice(symbols) for _ in range(k)]


def r_string(symbols, length):
    """Return random sequence from given symbols."""
    return ''.join(R.choice(symbols) for _ in range(length))


def r_symbols(size, symbols, length, used=None):
    """Return unique random from given symbols."""
    if length == 1 and not used:
        return R.sample(symbols, size)
    rset, used = set(), set(used or [])
    while len(rset) < size:
        s = r_string(symbols, R.randint(1, length))
        if s not in used:
            rset.add(s)
    return list(rset)


def r_consts(size, used=None):
    """Return size many unique constants."""
    return r_symbols(size, CONST_SYMBOLS, ARGS.constant_length, used)


<mask token>


def r_preds(size, used=None):
    """Return size many unique predicates."""
    return r_symbols(size, PRED_SYMBOLS, ARGS.predicate_length, used)


def write_p(pred):
    """Format single predicate tuple into string."""
    return PRED_T.format(pred[0], ARG_SEP.join(pred[1]))


<mask token>


def output(context, targets):
    """Print the context and given targets."""
    print('\n'.join([write_r(c) for c in context]))
    for t, v, s in targets:
        print(TARGET_T.format(write_r([t]), v, ','.join(map(str, s))))


def gen_task(context, targets, upreds):
    """Fill context with random preds and output program."""
    ctx = context.copy()
    for _ in range(ARGS.noise_size):
        rtask = R.randint(1, max(1, ARGS.task))
        rtask = 1 if ARGS.arity != 2 and rtask == 8 else rtask
        ntask = 'gen_task' + str(rtask)
        ctx.append(globals()[ntask](upreds))
    output(ctx, targets)


def add_pred(context, pred, upreds, uconsts, psuccess=0.0):
    """Fail a ground case predicate given context."""
    if R.random() < psuccess:
        context.append([pred])
    if R.random() < 0.5:
        args = pred[1].copy()
        args[R.randrange(len(args))] = r_consts(1, uconsts)[0]
        context.append([(pred[0], args)])
    if R.random() < 0.5:
        p = r_preds(1, upreds)[0]
        upreds.append(p)
        context.append([(p, pred[1])])


def gen_task1(upreds=None):
    """Ground instances only: p(a).q(c,b)."""
    preds = r_preds(2, upreds)
    args = r_consts(R.randint(1, ARGS.arity))
    rule = [(preds[0], args)]
    if upreds:
        return rule
    ctx = list()
    add_pred(ctx, rule[0], preds, args, 1.0)
    targets = [(rule[0], 1, [0])]
    args = r_consts(R.randint(1, 2))
    fpred = preds[1], args
    add_pred(ctx, fpred, preds, args)
    targets.append((fpred, 0, [-1]))
    gen_task(ctx, targets, preds)


def gen_task2(upreds=None):
    """Variablised facts only: p(X).q(X,Y)."""
    preds = r_preds(2, upreds)
    ctx, targets = list(), list()
    if R.random() < 0.5 and ARGS.arity > 1:
        v = r_vars(1)[0]
        rule = [(preds[0], [v, v])]
        if upreds:
            return rule
        ctx.append(rule)
        cs = r_consts(2)
        c = R.choice(cs)
        targets.append(((preds[0], [c, c]), 1, [0]))
        targets.append(((preds[0], cs), 0, [-1]))
    else:
        argc = R.randint(1, ARGS.arity)
        args = r_vars(argc)
        rule = [(preds[0], args)]
        if upreds:
            return rule
        ctx.append(rule)
        args = choices(r_consts(2), argc)
        targets.append(((preds[0], args), 1, [0]))
        targets.append(((preds[1], args), 0, [-1]))
    gen_task(ctx, targets, preds)


def nstep_deduction(steps, negation=False, upreds=None):
    """Generate nstep deduction programs given number of steps."""
    assert steps >= 1, 'Need at least 1 step deduction.'
    preds = r_preds(2 if upreds else 3 + steps, upreds)
    consts = r_consts(2)
    ctx, targets = list(), list()
    prefix = '-' if negation else ''
    if R.random() < 0.5 and ARGS.arity > 1:
        vs = r_vars(2)
        rule = [(preds[0], vs), (prefix + preds[1], vs[::-1])]
        if upreds:
            return rule
        ctx.append(rule)
        swapc = 1
        for j in range(steps - 1):
            vs = r_vars(2)
            toswap = R.random() < 0.5
            args = vs[::-1] if toswap else vs
            ctx.append([(preds[j + 1], vs), (preds[j + 2], args)])
            swapc += int(toswap)
        args = r_consts(2)
        add_pred(ctx, (preds[steps], args), preds, consts, 1.0)
        args = args if swapc % 2 == 0 else args[::-1]
        succ_supps = list(range(steps + 1))
        fail_supps = succ_supps[:-1] + [-1]
        targets.append(((preds[0], args), 1 - int(negation), succ_supps))
        targets.append(((preds[0], args[::-1]), int(negation), fail_supps))
        gen_task(ctx, targets, preds)
    else:
        argc = R.randint(1, ARGS.arity)
        vs = r_vars(argc)
        rule = [(preds[0], vs), (prefix + preds[1], vs)]
        if upreds:
            return rule
        ctx.append(rule)
        for j in range(steps - 1):
            vs = r_vars(argc)
            ctx.append([(preds[j + 1], vs), (preds[j + 2], vs)])
        args = choices(r_consts(2), argc)
        cctx = ctx.copy()
        spred = preds[steps], args
        add_pred(cctx, spred, preds, args, 1.0)
        succ_supps = list(range(steps + 1))
        targets = [((preds[0], args), 1 - int(negation), succ_supps)]
        gen_task(cctx, targets, preds)
        if R.random() < 0.5:
            p = r_preds(1, preds)[0]
            preds.append(p)
            add_pred(ctx, spred, preds, args, 1.0)
            ctx[0] = [(preds[0], vs), (prefix + p, vs)]
            supps = [0] + [-1] * steps
        else:
            add_pred(ctx, spred, preds, args)
            supps = succ_supps[:-1] + [-1]
        targets = [((preds[0], args), int(negation), supps)]
        gen_task(ctx, targets, preds)


def gen_task3(upreds=None):
    """Single step deduction: p(X):-q(X)."""
    return nstep_deduction(1, upreds=upreds)


def gen_task4(upreds=None):
    """Double step deduction: p(X):-q(X).q(X):-r(X)."""
    return nstep_deduction(2, upreds=upreds)


def gen_task5(upreds=None):
    """Triple step deduction."""
    return nstep_deduction(3, upreds=upreds)


def logical_and(negation=False, upreds=None):
    """Logical AND with optional negation: p(X):-q(X);r(X)."""
    preds = r_preds(3, upreds)
    argc = R.randint(1, ARGS.arity)
    vs = r_vars(argc)
    rule = [(preds[0], vs), (preds[1], vs[:1]), (preds[2], vs[1:] or vs)]
    if upreds:
        return rule
    ctx = [rule]
    args = choices(r_consts(2), argc)
    prem1 = preds[1], args[:1]
    prem2 = preds[2], args[1:] or args
    prems = [prem1, prem2]
    if negation:
        ridx = R.randrange(2)
        p, pargs = ctx[-1][ridx + 1]
        ctx[-1][ridx + 1] = '-' + p, pargs
        cctx = ctx.copy()
        add_pred(cctx, prems[ridx], preds, args)
        supps = [0, -1, -1] if ridx == 0 else [0, len(cctx), -1]
        cctx.append([prems[1 - ridx]])
        targets = [((preds[0], args), 1, supps)]
        gen_task(cctx, targets, preds)
        fidx = R.randrange(2)
        if ridx == fidx:
            supps = [0, len(ctx), -1]
            ctx.append([prems[ridx]])
            add_pred(ctx, prems[1 - ridx], preds, args, 0.8)
        else:
            supps = [0, -1, -1]
            add_pred(ctx, prems[1 - ridx], preds, args)
            add_pred(ctx, prems[ridx], preds, args)
        targets = [((preds[0], args), 0, supps)]
        gen_task(ctx, targets, preds)
    else:
        cctx = ctx.copy()
        add_pred(cctx, prems[0], preds, args, 1.0)
        supps = [0, 1, len(cctx)]
        add_pred(cctx, prems[1], preds, args, 1.0)
        targets = [((preds[0], args), 1, supps)]
        gen_task(cctx, targets, preds)
        fidx = R.randrange(2)
        add_pred(ctx, prems[fidx], preds, args)
        add_pred(ctx, prems[1 - fidx], preds, args, 0.8)
        targets = [((preds[0], args), 0, [0, -1, -1])]
        gen_task(ctx, targets, preds)


def gen_task6(upreds=None):
    """Logical AND: p(X):-q(X);r(X)."""
    return logical_and(upreds=upreds)


def logical_or(negation=False, upreds=None):
    """Logical OR with optional negation: p(X):-q(X).p(X):-r(X)."""
    preds = r_preds(3, upreds)
    argc = R.randint(1, ARGS.arity)
    vs = r_vars(argc)
    swap = R.random() < 0.5
    prefix = '-' if negation else ''
    rule = [(preds[0], vs), (prefix + preds[1], vs[::-1] if swap else vs)]
    if upreds:
        return rule
    ctx = list()
    ctx.append(rule)
    ctx.append([(preds[0], vs), (preds[2], vs)])
    args = r_consts(argc)
    ctx.append([(preds[0], args)])
    if swap and argc == 2:
        args = r_consts(argc, args)
        add_pred(ctx, (preds[1], args), preds, args, 1.0)
        args = args[::-1] if swap else args
        targets = [((preds[0], args), 1 - int(negation), [0, 3]), ((preds[0
            ], args[::-1]), int(negation), [0, -1])]
        gen_task(ctx, targets, preds)
    elif not negation and R.random() < 0.2:
        targets = [((preds[0], args), 1, [2, -1])]
        gen_task(ctx, targets, preds)
        del ctx[-1]
        targets = [((preds[0], args), 0, [-1, -1])]
        gen_task(ctx, targets, preds)
    else:
        prems = [(preds[i], r_consts(argc, args)) for i in range(1, 3)]
        sidx = R.randrange(2)
        cctx = ctx.copy()
        if negation and sidx == 0:
            supps = [0, -1]
            add_pred(cctx, prems[0], preds, prems[0][1])
            add_pred(cctx, prems[1], preds, prems[1][1], 0.2)
        else:
            supps = [sidx, len(cctx)]
            add_pred(cctx, prems[sidx], preds, prems[sidx][1], 1.0)
            add_pred(cctx, prems[1 - sidx], preds, prems[1 - sidx][1], 0.2)
        targets = [((preds[0], prems[sidx][1]), 1, supps)]
        gen_task(cctx, targets, preds)
        add_pred(ctx, prems[0], preds, prems[0][1], int(negation))
        add_pred(ctx, prems[1], preds, prems[1][1])
        targets = [((preds[0], prems[sidx][1]), 0, [-1, -1])]
        gen_task(ctx, targets, preds)


def gen_task7(upreds=None):
    """Logical OR: p(X):-q(X).p(X):-r(X)."""
    return logical_or(upreds=upreds)


<mask token>


def gen_task9(upreds=None):
    """Single step deduction with NBF: p(X):--q(X)."""
    return nstep_deduction(1, True, upreds)


def gen_task10(upreds=None):
    """Double step deduction with NBF: p(X):--q(X).q(X):-r(X)."""
    return nstep_deduction(2, True, upreds)


<mask token>


def gen_task12(upreds=None):
    """Logical OR with NBF: p(X):--q(X).p(X):-r(X)."""
    return logical_or(True, upreds)


def gen_task0():
    """Generate an ILP task example."""
    argc = 1
    goal = 'f'
    premise = 'b'
    ctx, targets = list(), list()
    args = r_consts(argc)
    ctx.append([(premise, args)])
    targets.append(((goal, args), 1))
    args = args.copy()
    args[R.randrange(len(args))] = r_consts(1, args)[0]
    preds = r_preds(3)
    ctx.append([(preds[0], args)])
    targets.append(((goal, args), 0))
    vs = r_vars(argc)
    ctx.append([(preds[1], vs), (preds[2], vs)])
    preds.extend([goal, premise])
    gen_task(ctx, targets, preds)


<mask token>

<mask token>


def choices(symbols, k):
    """Return k many symbols with replacement. Added in v3.6."""
    return [R.choice(symbols) for _ in range(k)]


def r_string(symbols, length):
    """Return random sequence from given symbols."""
    return ''.join(R.choice(symbols) for _ in range(length))


def r_symbols(size, symbols, length, used=None):
    """Return unique random from given symbols."""
    if length == 1 and not used:
        return R.sample(symbols, size)
    rset, used = set(), set(used or [])
    while len(rset) < size:
        s = r_string(symbols, R.randint(1, length))
        if s not in used:
            rset.add(s)
    return list(rset)


def r_consts(size, used=None):
    """Return size many unique constants."""
    return r_symbols(size, CONST_SYMBOLS, ARGS.constant_length, used)


<mask token>


def r_preds(size, used=None):
    """Return size many unique predicates."""
    return r_symbols(size, PRED_SYMBOLS, ARGS.predicate_length, used)


def write_p(pred):
    """Format single predicate tuple into string."""
    return PRED_T.format(pred[0], ARG_SEP.join(pred[1]))


def write_r(preds):
    """Convert rule predicate tuple into string."""
    head = write_p(preds[0])
    if len(preds) == 1:
        return FACT_T.format(head)
    return RULE_T.format(head, PRED_SEP.join([write_p(p) for p in preds[1:]]))


def output(context, targets):
    """Print the context and given targets."""
    print('\n'.join([write_r(c) for c in context]))
    for t, v, s in targets:
        print(TARGET_T.format(write_r([t]), v, ','.join(map(str, s))))


def gen_task(context, targets, upreds):
    """Fill context with random preds and output program."""
    ctx = context.copy()
    for _ in range(ARGS.noise_size):
        rtask = R.randint(1, max(1, ARGS.task))
        rtask = 1 if ARGS.arity != 2 and rtask == 8 else rtask
        ntask = 'gen_task' + str(rtask)
        ctx.append(globals()[ntask](upreds))
    output(ctx, targets)


def add_pred(context, pred, upreds, uconsts, psuccess=0.0):
    """Fail a ground case predicate given context."""
    if R.random() < psuccess:
        context.append([pred])
    if R.random() < 0.5:
        args = pred[1].copy()
        args[R.randrange(len(args))] = r_consts(1, uconsts)[0]
        context.append([(pred[0], args)])
    if R.random() < 0.5:
        p = r_preds(1, upreds)[0]
        upreds.append(p)
        context.append([(p, pred[1])])


def gen_task1(upreds=None):
    """Ground instances only: p(a).q(c,b)."""
    preds = r_preds(2, upreds)
    args = r_consts(R.randint(1, ARGS.arity))
    rule = [(preds[0], args)]
    if upreds:
        return rule
    ctx = list()
    add_pred(ctx, rule[0], preds, args, 1.0)
    targets = [(rule[0], 1, [0])]
    args = r_consts(R.randint(1, 2))
    fpred = preds[1], args
    add_pred(ctx, fpred, preds, args)
    targets.append((fpred, 0, [-1]))
    gen_task(ctx, targets, preds)


def gen_task2(upreds=None):
    """Variablised facts only: p(X).q(X,Y)."""
    preds = r_preds(2, upreds)
    ctx, targets = list(), list()
    if R.random() < 0.5 and ARGS.arity > 1:
        v = r_vars(1)[0]
        rule = [(preds[0], [v, v])]
        if upreds:
            return rule
        ctx.append(rule)
        cs = r_consts(2)
        c = R.choice(cs)
        targets.append(((preds[0], [c, c]), 1, [0]))
        targets.append(((preds[0], cs), 0, [-1]))
    else:
        argc = R.randint(1, ARGS.arity)
        args = r_vars(argc)
        rule = [(preds[0], args)]
        if upreds:
            return rule
        ctx.append(rule)
        args = choices(r_consts(2), argc)
        targets.append(((preds[0], args), 1, [0]))
        targets.append(((preds[1], args), 0, [-1]))
    gen_task(ctx, targets, preds)


def nstep_deduction(steps, negation=False, upreds=None):
    """Generate nstep deduction programs given number of steps."""
    assert steps >= 1, 'Need at least 1 step deduction.'
    preds = r_preds(2 if upreds else 3 + steps, upreds)
    consts = r_consts(2)
    ctx, targets = list(), list()
    prefix = '-' if negation else ''
    if R.random() < 0.5 and ARGS.arity > 1:
        vs = r_vars(2)
        rule = [(preds[0], vs), (prefix + preds[1], vs[::-1])]
        if upreds:
            return rule
        ctx.append(rule)
        swapc = 1
        for j in range(steps - 1):
            vs = r_vars(2)
            toswap = R.random() < 0.5
            args = vs[::-1] if toswap else vs
            ctx.append([(preds[j + 1], vs), (preds[j + 2], args)])
            swapc += int(toswap)
        args = r_consts(2)
        add_pred(ctx, (preds[steps], args), preds, consts, 1.0)
        args = args if swapc % 2 == 0 else args[::-1]
        succ_supps = list(range(steps + 1))
        fail_supps = succ_supps[:-1] + [-1]
        targets.append(((preds[0], args), 1 - int(negation), succ_supps))
        targets.append(((preds[0], args[::-1]), int(negation), fail_supps))
        gen_task(ctx, targets, preds)
    else:
        argc = R.randint(1, ARGS.arity)
        vs = r_vars(argc)
        rule = [(preds[0], vs), (prefix + preds[1], vs)]
        if upreds:
            return rule
        ctx.append(rule)
        for j in range(steps - 1):
            vs = r_vars(argc)
            ctx.append([(preds[j + 1], vs), (preds[j + 2], vs)])
        args = choices(r_consts(2), argc)
        cctx = ctx.copy()
        spred = preds[steps], args
        add_pred(cctx, spred, preds, args, 1.0)
        succ_supps = list(range(steps + 1))
        targets = [((preds[0], args), 1 - int(negation), succ_supps)]
        gen_task(cctx, targets, preds)
        if R.random() < 0.5:
            p = r_preds(1, preds)[0]
            preds.append(p)
            add_pred(ctx, spred, preds, args, 1.0)
            ctx[0] = [(preds[0], vs), (prefix + p, vs)]
            supps = [0] + [-1] * steps
        else:
            add_pred(ctx, spred, preds, args)
            supps = succ_supps[:-1] + [-1]
        targets = [((preds[0], args), int(negation), supps)]
        gen_task(ctx, targets, preds)


def gen_task3(upreds=None):
    """Single step deduction: p(X):-q(X)."""
    return nstep_deduction(1, upreds=upreds)


def gen_task4(upreds=None):
    """Double step deduction: p(X):-q(X).q(X):-r(X)."""
    return nstep_deduction(2, upreds=upreds)


def gen_task5(upreds=None):
    """Triple step deduction."""
    return nstep_deduction(3, upreds=upreds)


def logical_and(negation=False, upreds=None):
    """Logical AND with optional negation: p(X):-q(X);r(X)."""
    preds = r_preds(3, upreds)
    argc = R.randint(1, ARGS.arity)
    vs = r_vars(argc)
    rule = [(preds[0], vs), (preds[1], vs[:1]), (preds[2], vs[1:] or vs)]
    if upreds:
        return rule
    ctx = [rule]
    args = choices(r_consts(2), argc)
    prem1 = preds[1], args[:1]
    prem2 = preds[2], args[1:] or args
    prems = [prem1, prem2]
    if negation:
        ridx = R.randrange(2)
        p, pargs = ctx[-1][ridx + 1]
        ctx[-1][ridx + 1] = '-' + p, pargs
        cctx = ctx.copy()
        add_pred(cctx, prems[ridx], preds, args)
        supps = [0, -1, -1] if ridx == 0 else [0, len(cctx), -1]
        cctx.append([prems[1 - ridx]])
        targets = [((preds[0], args), 1, supps)]
        gen_task(cctx, targets, preds)
        fidx = R.randrange(2)
        if ridx == fidx:
            supps = [0, len(ctx), -1]
            ctx.append([prems[ridx]])
            add_pred(ctx, prems[1 - ridx], preds, args, 0.8)
        else:
            supps = [0, -1, -1]
            add_pred(ctx, prems[1 - ridx], preds, args)
            add_pred(ctx, prems[ridx], preds, args)
        targets = [((preds[0], args), 0, supps)]
        gen_task(ctx, targets, preds)
    else:
        cctx = ctx.copy()
        add_pred(cctx, prems[0], preds, args, 1.0)
        supps = [0, 1, len(cctx)]
        add_pred(cctx, prems[1], preds, args, 1.0)
        targets = [((preds[0], args), 1, supps)]
        gen_task(cctx, targets, preds)
        fidx = R.randrange(2)
        add_pred(ctx, prems[fidx], preds, args)
        add_pred(ctx, prems[1 - fidx], preds, args, 0.8)
        targets = [((preds[0], args), 0, [0, -1, -1])]
        gen_task(ctx, targets, preds)


def gen_task6(upreds=None):
    """Logical AND: p(X):-q(X);r(X)."""
    return logical_and(upreds=upreds)


def logical_or(negation=False, upreds=None):
    """Logical OR with optional negation: p(X):-q(X).p(X):-r(X)."""
    preds = r_preds(3, upreds)
    argc = R.randint(1, ARGS.arity)
    vs = r_vars(argc)
    swap = R.random() < 0.5
    prefix = '-' if negation else ''
    rule = [(preds[0], vs), (prefix + preds[1], vs[::-1] if swap else vs)]
    if upreds:
        return rule
    ctx = list()
    ctx.append(rule)
    ctx.append([(preds[0], vs), (preds[2], vs)])
    args = r_consts(argc)
    ctx.append([(preds[0], args)])
    if swap and argc == 2:
        args = r_consts(argc, args)
        add_pred(ctx, (preds[1], args), preds, args, 1.0)
        args = args[::-1] if swap else args
        targets = [((preds[0], args), 1 - int(negation), [0, 3]), ((preds[0
            ], args[::-1]), int(negation), [0, -1])]
        gen_task(ctx, targets, preds)
    elif not negation and R.random() < 0.2:
        targets = [((preds[0], args), 1, [2, -1])]
        gen_task(ctx, targets, preds)
        del ctx[-1]
        targets = [((preds[0], args), 0, [-1, -1])]
        gen_task(ctx, targets, preds)
    else:
        prems = [(preds[i], r_consts(argc, args)) for i in range(1, 3)]
        sidx = R.randrange(2)
        cctx = ctx.copy()
        if negation and sidx == 0:
            supps = [0, -1]
            add_pred(cctx, prems[0], preds, prems[0][1])
            add_pred(cctx, prems[1], preds, prems[1][1], 0.2)
        else:
            supps = [sidx, len(cctx)]
            add_pred(cctx, prems[sidx], preds, prems[sidx][1], 1.0)
            add_pred(cctx, prems[1 - sidx], preds, prems[1 - sidx][1], 0.2)
        targets = [((preds[0], prems[sidx][1]), 1, supps)]
        gen_task(cctx, targets, preds)
        add_pred(ctx, prems[0], preds, prems[0][1], int(negation))
        add_pred(ctx, prems[1], preds, prems[1][1])
        targets = [((preds[0], prems[sidx][1]), 0, [-1, -1])]
        gen_task(ctx, targets, preds)


def gen_task7(upreds=None):
    """Logical OR: p(X):-q(X).p(X):-r(X)."""
    return logical_or(upreds=upreds)


<mask token>


def gen_task9(upreds=None):
    """Single step deduction with NBF: p(X):--q(X)."""
    return nstep_deduction(1, True, upreds)


def gen_task10(upreds=None):
    """Double step deduction with NBF: p(X):--q(X).q(X):-r(X)."""
    return nstep_deduction(2, True, upreds)


<mask token>


def gen_task12(upreds=None):
    """Logical OR with NBF: p(X):--q(X).p(X):-r(X)."""
    return logical_or(True, upreds)


def gen_task0():
    """Generate an ILP task example."""
    argc = 1
    goal = 'f'
    premise = 'b'
    ctx, targets = list(), list()
    args = r_consts(argc)
    ctx.append([(premise, args)])
    targets.append(((goal, args), 1))
    args = args.copy()
    args[R.randrange(len(args))] = r_consts(1, args)[0]
    preds = r_preds(3)
    ctx.append([(preds[0], args)])
    targets.append(((goal, args), 0))
    vs = r_vars(argc)
    ctx.append([(preds[1], vs), (preds[2], vs)])
    preds.extend([goal, premise])
    gen_task(ctx, targets, preds)


<mask token>

<mask token>
import argparse
import random as R
CONST_SYMBOLS = 'abcdefghijklmnopqrstuvwxyz'
VAR_SYMBOLS = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
PRED_SYMBOLS = 'abcdefghijklmnopqrstuvwxyz'
EXTRA_SYMBOLS = '-,()'
CHARS = sorted(list(set(CONST_SYMBOLS + VAR_SYMBOLS + PRED_SYMBOLS +
    EXTRA_SYMBOLS)))
CHAR_IDX = dict((c, i + 1) for i, c in enumerate(CHARS))
IDX_CHAR = [0]
IDX_CHAR.extend(CHARS)
FACT_T = '{}.'
RULE_T = '{}:-{}.'
PRED_T = '{}({})'
ARG_SEP = ','
PRED_SEP = ';'
TARGET_T = '? {} {} {}'


def choices(symbols, k):
    """Return k many symbols with replacement. Added in v3.6."""
    return [R.choice(symbols) for _ in range(k)]


def r_string(symbols, length):
    """Return random sequence from given symbols."""
    return ''.join(R.choice(symbols) for _ in range(length))


def r_symbols(size, symbols, length, used=None):
    """Return unique random from given symbols."""
    if length == 1 and not used:
        return R.sample(symbols, size)
    rset, used = set(), set(used or [])
    while len(rset) < size:
        s = r_string(symbols, R.randint(1, length))
        if s not in used:
            rset.add(s)
    return list(rset)


def r_consts(size, used=None):
    """Return size many unique constants."""
    return r_symbols(size, CONST_SYMBOLS, ARGS.constant_length, used)


def r_vars(size, used=None):
    """Return size many unique variables."""
    return r_symbols(size, VAR_SYMBOLS, ARGS.variable_length, used)


def r_preds(size, used=None):
    """Return size many unique predicates."""
    return r_symbols(size, PRED_SYMBOLS, ARGS.predicate_length, used)


def write_p(pred):
    """Format single predicate tuple into string."""
    return PRED_T.format(pred[0], ARG_SEP.join(pred[1]))


def write_r(preds):
    """Convert rule predicate tuple into string."""
    head = write_p(preds[0])
    if len(preds) == 1:
        return FACT_T.format(head)
    return RULE_T.format(head, PRED_SEP.join([write_p(p) for p in preds[1:]]))


def output(context, targets):
    """Print the context and given targets."""
    print('\n'.join([write_r(c) for c in context]))
    for t, v, s in targets:
        print(TARGET_T.format(write_r([t]), v, ','.join(map(str, s))))


def gen_task(context, targets, upreds):
    """Fill context with random preds and output program."""
    ctx = context.copy()
    for _ in range(ARGS.noise_size):
        rtask = R.randint(1, max(1, ARGS.task))
        rtask = 1 if ARGS.arity != 2 and rtask == 8 else rtask
        ntask = 'gen_task' + str(rtask)
        ctx.append(globals()[ntask](upreds))
    output(ctx, targets)


def add_pred(context, pred, upreds, uconsts, psuccess=0.0):
    """Fail a ground case predicate given context."""
    if R.random() < psuccess:
        context.append([pred])
    if R.random() < 0.5:
        args = pred[1].copy()
        args[R.randrange(len(args))] = r_consts(1, uconsts)[0]
        context.append([(pred[0], args)])
    if R.random() < 0.5:
        p = r_preds(1, upreds)[0]
        upreds.append(p)
        context.append([(p, pred[1])])


def gen_task1(upreds=None):
    """Ground instances only: p(a).q(c,b)."""
    preds = r_preds(2, upreds)
    args = r_consts(R.randint(1, ARGS.arity))
    rule = [(preds[0], args)]
    if upreds:
        return rule
    ctx = list()
    add_pred(ctx, rule[0], preds, args, 1.0)
    targets = [(rule[0], 1, [0])]
    args = r_consts(R.randint(1, 2))
    fpred = preds[1], args
    add_pred(ctx, fpred, preds, args)
    targets.append((fpred, 0, [-1]))
    gen_task(ctx, targets, preds)


def gen_task2(upreds=None):
    """Variablised facts only: p(X).q(X,Y)."""
    preds = r_preds(2, upreds)
    ctx, targets = list(), list()
    if R.random() < 0.5 and ARGS.arity > 1:
        v = r_vars(1)[0]
        rule = [(preds[0], [v, v])]
        if upreds:
            return rule
        ctx.append(rule)
        cs = r_consts(2)
        c = R.choice(cs)
        targets.append(((preds[0], [c, c]), 1, [0]))
        targets.append(((preds[0], cs), 0, [-1]))
    else:
        argc = R.randint(1, ARGS.arity)
        args = r_vars(argc)
        rule = [(preds[0], args)]
        if upreds:
            return rule
        ctx.append(rule)
        args = choices(r_consts(2), argc)
        targets.append(((preds[0], args), 1, [0]))
        targets.append(((preds[1], args), 0, [-1]))
    gen_task(ctx, targets, preds)


def nstep_deduction(steps, negation=False, upreds=None):
    """Generate nstep deduction programs given number of steps."""
    assert steps >= 1, 'Need at least 1 step deduction.'
    preds = r_preds(2 if upreds else 3 + steps, upreds)
    consts = r_consts(2)
    ctx, targets = list(), list()
    prefix = '-' if negation else ''
    if R.random() < 0.5 and ARGS.arity > 1:
        vs = r_vars(2)
        rule = [(preds[0], vs), (prefix + preds[1], vs[::-1])]
        if upreds:
            return rule
        ctx.append(rule)
        swapc = 1
        for j in range(steps - 1):
            vs = r_vars(2)
            toswap = R.random() < 0.5
            args = vs[::-1] if toswap else vs
            ctx.append([(preds[j + 1], vs), (preds[j + 2], args)])
            swapc += int(toswap)
        args = r_consts(2)
        add_pred(ctx, (preds[steps], args), preds, consts, 1.0)
        args = args if swapc % 2 == 0 else args[::-1]
        succ_supps = list(range(steps + 1))
        fail_supps = succ_supps[:-1] + [-1]
        targets.append(((preds[0], args), 1 - int(negation), succ_supps))
        targets.append(((preds[0], args[::-1]), int(negation), fail_supps))
        gen_task(ctx, targets, preds)
    else:
        argc = R.randint(1, ARGS.arity)
        vs = r_vars(argc)
        rule = [(preds[0], vs), (prefix + preds[1], vs)]
        if upreds:
            return rule
        ctx.append(rule)
        for j in range(steps - 1):
            vs = r_vars(argc)
            ctx.append([(preds[j + 1], vs), (preds[j + 2], vs)])
        args = choices(r_consts(2), argc)
        cctx = ctx.copy()
        spred = preds[steps], args
        add_pred(cctx, spred, preds, args, 1.0)
        succ_supps = list(range(steps + 1))
        targets = [((preds[0], args), 1 - int(negation), succ_supps)]
        gen_task(cctx, targets, preds)
        if R.random() < 0.5:
            p = r_preds(1, preds)[0]
            preds.append(p)
            add_pred(ctx, spred, preds, args, 1.0)
            ctx[0] = [(preds[0], vs), (prefix + p, vs)]
            supps = [0] + [-1] * steps
        else:
            add_pred(ctx, spred, preds, args)
            supps = succ_supps[:-1] + [-1]
        targets = [((preds[0], args), int(negation), supps)]
        gen_task(ctx, targets, preds)


def gen_task3(upreds=None):
    """Single step deduction: p(X):-q(X)."""
    return nstep_deduction(1, upreds=upreds)


def gen_task4(upreds=None):
    """Double step deduction: p(X):-q(X).q(X):-r(X)."""
    return nstep_deduction(2, upreds=upreds)


def gen_task5(upreds=None):
    """Triple step deduction."""
    return nstep_deduction(3, upreds=upreds)


def logical_and(negation=False, upreds=None):
    """Logical AND with optional negation: p(X):-q(X);r(X)."""
    preds = r_preds(3, upreds)
    argc = R.randint(1, ARGS.arity)
    vs = r_vars(argc)
    rule = [(preds[0], vs), (preds[1], vs[:1]), (preds[2], vs[1:] or vs)]
    if upreds:
        return rule
    ctx = [rule]
    args = choices(r_consts(2), argc)
    prem1 = preds[1], args[:1]
    prem2 = preds[2], args[1:] or args
    prems = [prem1, prem2]
    if negation:
        ridx = R.randrange(2)
        p, pargs = ctx[-1][ridx + 1]
        ctx[-1][ridx + 1] = '-' + p, pargs
        cctx = ctx.copy()
        add_pred(cctx, prems[ridx], preds, args)
        supps = [0, -1, -1] if ridx == 0 else [0, len(cctx), -1]
        cctx.append([prems[1 - ridx]])
        targets = [((preds[0], args), 1, supps)]
        gen_task(cctx, targets, preds)
        fidx = R.randrange(2)
        if ridx == fidx:
            supps = [0, len(ctx), -1]
            ctx.append([prems[ridx]])
            add_pred(ctx, prems[1 - ridx], preds, args, 0.8)
        else:
            supps = [0, -1, -1]
            add_pred(ctx, prems[1 - ridx], preds, args)
            add_pred(ctx, prems[ridx], preds, args)
        targets = [((preds[0], args), 0, supps)]
        gen_task(ctx, targets, preds)
    else:
        cctx = ctx.copy()
        add_pred(cctx, prems[0], preds, args, 1.0)
        supps = [0, 1, len(cctx)]
        add_pred(cctx, prems[1], preds, args, 1.0)
        targets = [((preds[0], args), 1, supps)]
        gen_task(cctx, targets, preds)
        fidx = R.randrange(2)
        add_pred(ctx, prems[fidx], preds, args)
        add_pred(ctx, prems[1 - fidx], preds, args, 0.8)
        targets = [((preds[0], args), 0, [0, -1, -1])]
        gen_task(ctx, targets, preds)


def gen_task6(upreds=None):
    """Logical AND: p(X):-q(X);r(X)."""
    return logical_and(upreds=upreds)


def logical_or(negation=False, upreds=None):
    """Logical OR with optional negation: p(X):-q(X).p(X):-r(X)."""
    preds = r_preds(3, upreds)
    argc = R.randint(1, ARGS.arity)
    vs = r_vars(argc)
    swap = R.random() < 0.5
    prefix = '-' if negation else ''
    rule = [(preds[0], vs), (prefix + preds[1], vs[::-1] if swap else vs)]
    if upreds:
        return rule
    ctx = list()
    ctx.append(rule)
    ctx.append([(preds[0], vs), (preds[2], vs)])
    args = r_consts(argc)
    ctx.append([(preds[0], args)])
    if swap and argc == 2:
        args = r_consts(argc, args)
        add_pred(ctx, (preds[1], args), preds, args, 1.0)
        args = args[::-1] if swap else args
        targets = [((preds[0], args), 1 - int(negation), [0, 3]), ((preds[0
            ], args[::-1]), int(negation), [0, -1])]
        gen_task(ctx, targets, preds)
    elif not negation and R.random() < 0.2:
        targets = [((preds[0], args), 1, [2, -1])]
        gen_task(ctx, targets, preds)
        del ctx[-1]
        targets = [((preds[0], args), 0, [-1, -1])]
        gen_task(ctx, targets, preds)
    else:
        prems = [(preds[i], r_consts(argc, args)) for i in range(1, 3)]
        sidx = R.randrange(2)
        cctx = ctx.copy()
        if negation and sidx == 0:
            supps = [0, -1]
            add_pred(cctx, prems[0], preds, prems[0][1])
            add_pred(cctx, prems[1], preds, prems[1][1], 0.2)
        else:
            supps = [sidx, len(cctx)]
            add_pred(cctx, prems[sidx], preds, prems[sidx][1], 1.0)
            add_pred(cctx, prems[1 - sidx], preds, prems[1 - sidx][1], 0.2)
        targets = [((preds[0], prems[sidx][1]), 1, supps)]
        gen_task(cctx, targets, preds)
        add_pred(ctx, prems[0], preds, prems[0][1], int(negation))
        add_pred(ctx, prems[1], preds, prems[1][1])
        targets = [((preds[0], prems[sidx][1]), 0, [-1, -1])]
        gen_task(ctx, targets, preds)


def gen_task7(upreds=None):
    """Logical OR: p(X):-q(X).p(X):-r(X)."""
    return logical_or(upreds=upreds)


def gen_task8(upreds=None):
    """Transitive case: p(X,Y):-q(X,Z);r(Z,Y)."""
    assert ARGS.arity == 2, 'Transitive task arity must be 2.'
    preds = r_preds(3, upreds)
    vs = r_vars(3)
    rule = [(preds[0], [vs[0], vs[2]]), (preds[1], vs[:2]), (preds[2], vs[1:])]
    if upreds:
        return rule
    ctx = [rule]
    args = r_consts(3)
    add_pred(ctx, (preds[1], args[:2]), preds, args, 1.0)
    succ_supps = [0, 1, len(ctx), -1]
    add_pred(ctx, (preds[2], args[1:]), preds, args, 1.0)
    argso = r_consts(3)
    argso.insert(R.randint(1, 2), r_consts(1, argso)[0])
    add_pred(ctx, (preds[1], argso[:2]), preds, argso, 0.5)
    add_pred(ctx, (preds[2], argso[2:]), preds, argso, 0.5)
    targets = [((preds[0], [args[0], args[2]]), 1, succ_supps), ((preds[0],
        [argso[0], argso[3]]), 0, [0, -1, -1, -1])]
    gen_task(ctx, targets, preds)


def gen_task9(upreds=None):
    """Single step deduction with NBF: p(X):--q(X)."""
    return nstep_deduction(1, True, upreds)


def gen_task10(upreds=None):
    """Double step deduction with NBF: p(X):--q(X).q(X):-r(X)."""
    return nstep_deduction(2, True, upreds)


def gen_task11(upreds=None):
    """Logical AND with NBF: p(X):-q(X);-r(X)."""
    return logical_and(True, upreds)


def gen_task12(upreds=None):
    """Logical OR with NBF: p(X):--q(X).p(X):-r(X)."""
    return logical_or(True, upreds)


def gen_task0():
    """Generate an ILP task example."""
    argc = 1
    goal = 'f'
    premise = 'b'
    ctx, targets = list(), list()
    args = r_consts(argc)
    ctx.append([(premise, args)])
    targets.append(((goal, args), 1))
    args = args.copy()
    args[R.randrange(len(args))] = r_consts(1, args)[0]
    preds = r_preds(3)
    ctx.append([(preds[0], args)])
    targets.append(((goal, args), 0))
    vs = r_vars(argc)
    ctx.append([(preds[1], vs), (preds[2], vs)])
    preds.extend([goal, premise])
    gen_task(ctx, targets, preds)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate logic program data.'
        )
    parser.add_argument('-t', '--task', default=1, type=int, help=
        'The task to generate.')
    parser.add_argument('-s', '--size', default=1, type=int, help=
        'Number of programs to generate.')
    parser.add_argument('-ns', '--noise_size', default=2, type=int, help=
        'Size of added noise rules.')
    parser.add_argument('-cl', '--constant_length', default=1, type=int,
        help='Length of constants.')
    parser.add_argument('-vl', '--variable_length', default=1, type=int,
        help='Length of variables.')
    parser.add_argument('-pl', '--predicate_length', default=1, type=int,
        help='Length of predicates.')
    parser.add_argument('-ar', '--arity', default=2, type=int, help='Arity.')
    parser.add_argument('--nstep', type=int, help=
        'Generate nstep deduction programs.')
    ARGS = parser.parse_args()
    task = 'gen_task' + str(ARGS.task)
    for _ in range(ARGS.size):
        if ARGS.nstep:
            nstep_deduction(ARGS.nstep)
        else:
            globals()[task]()

"""Data generation script for logic programs."""
import argparse
import random as R

# Symbol Pool
CONST_SYMBOLS = "abcdefghijklmnopqrstuvwxyz"
VAR_SYMBOLS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
PRED_SYMBOLS = "abcdefghijklmnopqrstuvwxyz"
EXTRA_SYMBOLS = "-,()"

CHARS = sorted(list(set(CONST_SYMBOLS+VAR_SYMBOLS+PRED_SYMBOLS+EXTRA_SYMBOLS)))
# Reserve 0 for padding
CHAR_IDX = dict((c, i+1) for i, c in enumerate(CHARS))
IDX_CHAR = [0]
IDX_CHAR.extend(CHARS)

# Predicate Templates
FACT_T = "{}."
RULE_T = "{}:-{}."
PRED_T = "{}({})"
ARG_SEP = ','
PRED_SEP = ';'
TARGET_T = "? {} {} {}"

def choices(symbols, k):
  """Return k many symbols with replacement. Added in v3.6."""
  return [R.choice(symbols) for _ in range(k)]

def r_string(symbols, length):
  """Return random sequence from given symbols."""
  return ''.join(R.choice(symbols)
                 for _ in range(length))

def r_symbols(size, symbols, length, used=None):
  """Return unique random from given symbols."""
  if length == 1 and not used:
    return R.sample(symbols, size)
  rset, used = set(), set(used or [])
  while len(rset) < size:
    s = r_string(symbols, R.randint(1, length))
    if s not in used:
      rset.add(s)
  return list(rset)

def r_consts(size, used=None):
  """Return size many unique constants."""
  return r_symbols(size, CONST_SYMBOLS, ARGS.constant_length, used)

def r_vars(size, used=None):
  """Return size many unique variables."""
  return r_symbols(size, VAR_SYMBOLS, ARGS.variable_length, used)

def r_preds(size, used=None):
  """Return size many unique predicates."""
  return r_symbols(size, PRED_SYMBOLS, ARGS.predicate_length, used)

def write_p(pred):
  """Format single predicate tuple into string."""
  return PRED_T.format(pred[0], ARG_SEP.join(pred[1]))

def write_r(preds):
  """Convert rule predicate tuple into string."""
  head = write_p(preds[0])
  # Is it just a fact
  if len(preds) == 1:
    return FACT_T.format(head)
  # We have a rule
  return RULE_T.format(head, PRED_SEP.join([write_p(p) for p in preds[1:]]))

def output(context, targets):
  """Print the context and given targets."""
  # context: [[('p', ['a', 'b'])], ...]
  # targets: [(('p', ['a', 'b']), 1, [0,1,2]), ...]
  print('\n'.join([write_r(c) for c in context]))
  for t, v, s in targets:
    print(TARGET_T.format(write_r([t]), v, ','.join(map(str, s))))

def gen_task(context, targets, upreds):
  """Fill context with random preds and output program."""
  # Fill with random rules up to certain task
  ctx = context.copy() # Don't modify the original context
  for _ in range(ARGS.noise_size):
    rtask = R.randint(1, max(1, ARGS.task))
    rtask = 1 if ARGS.arity != 2 and rtask == 8 else rtask
    ntask = "gen_task" + str(rtask)
    ctx.append(globals()[ntask](upreds))
  output(ctx, targets)

def add_pred(context, pred, upreds, uconsts, psuccess=0.0):
  """Fail a ground case predicate given context."""
  # Maybe succeed by adding to context
  if R.random() < psuccess:
    context.append([pred])
  if R.random() < 0.5:
    # The constant doesn't match
    args = pred[1].copy()
    args[R.randrange(len(args))] = r_consts(1, uconsts)[0]
    context.append([(pred[0], args)])
  if R.random() < 0.5:
    # The predicate doesn't match
    p = r_preds(1, upreds)[0]
    upreds.append(p)
    context.append([(p, pred[1])])
  # The predicate doesn't appear at all

def gen_task1(upreds=None):
  """Ground instances only: p(a).q(c,b)."""
  # One or two argument predicate
  preds = r_preds(2, upreds)
  args = r_consts(R.randint(1, ARGS.arity))
  rule = [(preds[0], args)]
  if upreds:
    return rule
  ctx = list()
  add_pred(ctx, rule[0], preds, args, 1.0)
  # Successful case when query appears in context
  targets = [(rule[0], 1, [0])]
  # Fail case
  args = r_consts(R.randint(1, 2))
  fpred = (preds[1], args)
  add_pred(ctx, fpred, preds, args)
  targets.append((fpred, 0, [-1]))
  gen_task(ctx, targets, preds)

def gen_task2(upreds=None):
  """Variablised facts only: p(X).q(X,Y)."""
  preds = r_preds(2, upreds)
  ctx, targets = list(), list()
  if R.random() < 0.5 and ARGS.arity > 1:
    # Double variable same argument
    v = r_vars(1)[0]
    rule = [(preds[0], [v, v])]
    if upreds:
      return rule
    ctx.append(rule)
    # Successful double variable grounding
    cs = r_consts(2)
    c = R.choice(cs)
    targets.append(((preds[0], [c, c]), 1, [0]))
    # Fail on non-unique variable grounding
    targets.append(((preds[0], cs), 0, [-1]))
  else:
    # Double variable different argument
    # Single variable argument
    argc = R.randint(1, ARGS.arity)
    args = r_vars(argc)
    rule = [(preds[0], args)]
    if upreds:
      return rule
    ctx.append(rule)
    # Successful unique argument grounding
    args = choices(r_consts(2), argc)
    targets.append(((preds[0], args), 1, [0]))
    # Fail on out of context predicate with same arguments
    targets.append(((preds[1], args), 0, [-1]))
  gen_task(ctx, targets, preds)

def nstep_deduction(steps, negation=False, upreds=None):
  """Generate nstep deduction programs given number of steps."""
  assert steps >= 1, "Need at least 1 step deduction."
  preds = r_preds(2 if upreds else 3+steps, upreds)
  consts = r_consts(2)
  ctx, targets = list(), list()
  prefix = '-' if negation else ''
  if R.random() < 0.5 and ARGS.arity > 1:
    # Double variable swap deduction rules
    vs = r_vars(2)
    rule = [(preds[0], vs), (prefix+preds[1], vs[::-1])]
    if upreds:
      return rule
    ctx.append(rule)
    # Add the n steps
    swapc = 1
    for j in range(steps-1):
      vs = r_vars(2)
      toswap = R.random() < 0.5 # Do we swap again?
      args = vs[::-1] if toswap else vs
      ctx.append([(preds[j+1], vs), (preds[j+2], args)])
      swapc += int(toswap)
    # Add the ground case
    args = r_consts(2)
    add_pred(ctx, (preds[steps], args), preds, consts, 1.0)
    args = args if swapc % 2 == 0 else args[::-1]
    succ_supps = list(range(steps+1))
    fail_supps = succ_supps[:-1] + [-1]
    targets.append(((preds[0], args), 1-int(negation), succ_supps))
    targets.append(((preds[0], args[::-1]), int(negation), fail_supps))
    gen_task(ctx, targets, preds)
  else:
    # Double variable non-swap deduction rules
    # Single variable deduction rules
    argc = R.randint(1, ARGS.arity)
    vs = r_vars(argc)
    rule = [(preds[0], vs), (prefix+preds[1], vs)]
    if upreds:
      return rule
    ctx.append(rule)
    # Add the n steps
    for j in range(steps-1):
      vs = r_vars(argc)
      ctx.append([(preds[j+1], vs), (preds[j+2], vs)])
    args = choices(r_consts(2), argc)
    # Add the ground case
    cctx = ctx.copy()
    spred = (preds[steps], args)
    add_pred(cctx, spred, preds, args, 1.0)
    succ_supps = list(range(steps+1))
    targets = [((preds[0], args), 1-int(negation), succ_supps)]
    gen_task(cctx, targets, preds)
    # Add failure case
    if R.random() < 0.5:
      # Fail on broken chain
      p = r_preds(1, preds)[0]
      preds.append(p)
      add_pred(ctx, spred, preds, args, 1.0)
      ctx[0] = [(preds[0], vs), (prefix+p, vs)]
      supps = [0] + [-1]*steps
    else:
      # Fail on last ground case
      add_pred(ctx, spred, preds, args)
      supps = succ_supps[:-1] + [-1]
    targets = [((preds[0], args), int(negation), supps)]
    gen_task(ctx, targets, preds)

def gen_task3(upreds=None):
  """Single step deduction: p(X):-q(X)."""
  return nstep_deduction(1, upreds=upreds)

def gen_task4(upreds=None):
  """Double step deduction: p(X):-q(X).q(X):-r(X)."""
  return nstep_deduction(2, upreds=upreds)

def gen_task5(upreds=None):
  """Triple step deduction."""
  return nstep_deduction(3, upreds=upreds)

def logical_and(negation=False, upreds=None):
  """Logical AND with optional negation: p(X):-q(X);r(X)."""
  preds = r_preds(3, upreds)
  argc = R.randint(1, ARGS.arity)
  # Double variable AND with different vars
  # Single variable AND
  vs = r_vars(argc)
  rule = [(preds[0], vs),
          (preds[1], vs[:1]),
          (preds[2], vs[1:] or vs)]
  if upreds:
    return rule
  ctx = [rule]
  # Create the ground arguments
  args = choices(r_consts(2), argc)
  prem1 = (preds[1], args[:1])
  prem2 = (preds[2], args[1:] or args)
  prems = [prem1, prem2]
  if negation:
    # Add negation to random predicate in body
    ridx = R.randrange(2)
    p, pargs = ctx[-1][ridx+1]
    ctx[-1][ridx+1] = ('-' + p, pargs)
    # Successful case when negation fails
    cctx = ctx.copy()
    add_pred(cctx, prems[ridx], preds, args)
    supps = [0, -1, -1] if ridx == 0 else [0, len(cctx), -1]
    cctx.append([prems[1-ridx]])
    targets = [((preds[0], args), 1, supps)]
    gen_task(cctx, targets, preds)
    # Fail one premise randomly
    fidx = R.randrange(2)
    if ridx == fidx:
      # To fail negation add ground instance
      supps = [0, len(ctx), -1]
      ctx.append([prems[ridx]])
      # Succeed other with some probability
      add_pred(ctx, prems[1-ridx], preds, args, 0.8)
    else:
      # Fail non-negated premise
      supps = [0, -1, -1]
      add_pred(ctx, prems[1-ridx], preds, args)
      # Still succeed negation
      add_pred(ctx, prems[ridx], preds, args)
    targets = [((preds[0], args), 0, supps)]
    gen_task(ctx, targets, preds)
  else:
    # Create successful context
    cctx = ctx.copy()
    add_pred(cctx, prems[0], preds, args, 1.0)
    supps = [0, 1, len(cctx)]
    add_pred(cctx, prems[1], preds, args, 1.0)
    targets = [((preds[0], args), 1, supps)]
    gen_task(cctx, targets, preds)
    # Fail one premise randomly
    fidx = R.randrange(2)
    add_pred(ctx, prems[fidx], preds, args)
    # Succeed the other with some probability
    add_pred(ctx, prems[1-fidx], preds, args, 0.8)
    targets = [((preds[0], args), 0, [0, -1, -1])]
    gen_task(ctx, targets, preds)

def gen_task6(upreds=None):
  """Logical AND: p(X):-q(X);r(X)."""
  return logical_and(upreds=upreds)

def logical_or(negation=False, upreds=None):
  """Logical OR with optional negation: p(X):-q(X).p(X):-r(X)."""
  preds = r_preds(3, upreds)
  # Double or single variable OR
  argc = R.randint(1, ARGS.arity)
  vs = r_vars(argc)
  swap = R.random() < 0.5
  prefix = '-' if negation else ''
  rule = [(preds[0], vs), (prefix + preds[1], vs[::-1] if swap else vs)]
  if upreds:
    return rule
  ctx = list()
  ctx.append(rule)
  # Add the extra branching rules
  ctx.append([(preds[0], vs), (preds[2], vs)])
  args = r_consts(argc)
  ctx.append([(preds[0], args)])
  if swap and argc == 2:
    args = r_consts(argc, args)
    add_pred(ctx, (preds[1], args), preds, args, 1.0)
    args = args[::-1] if swap else args
    targets = [((preds[0], args), 1-int(negation), [0, 3]),
               ((preds[0], args[::-1]), int(negation), [0, -1])]
    gen_task(ctx, targets, preds)
  elif not negation and R.random() < 0.2:
    # Sneaky shorcut case
    targets = [((preds[0], args), 1, [2, -1])]
    gen_task(ctx, targets, preds)
    del ctx[-1]
    targets = [((preds[0], args), 0, [-1, -1])]
    gen_task(ctx, targets, preds)
  else:
    # Succeed either from them
    prems = [(preds[i], r_consts(argc, args)) for i in range(1, 3)]
    sidx = R.randrange(2)
    cctx = ctx.copy()
    if negation and sidx == 0:
      # Succeed by failing negation
      supps = [0, -1]
      add_pred(cctx, prems[0], preds, prems[0][1])
      # Possibly succeed other prem
      add_pred(cctx, prems[1], preds, prems[1][1], 0.2)
    else:
      # Succeed by adding ground case
      supps = [sidx, len(cctx)]
      add_pred(cctx, prems[sidx], preds, prems[sidx][1], 1.0)
      # Possibly succeed other prem
      add_pred(cctx, prems[1-sidx], preds, prems[1-sidx][1], 0.2)
    targets = [((preds[0], prems[sidx][1]), 1, supps)]
    gen_task(cctx, targets, preds)
    # Fail both
    add_pred(ctx, prems[0], preds, prems[0][1], int(negation))
    add_pred(ctx, prems[1], preds, prems[1][1])
    targets = [((preds[0], prems[sidx][1]), 0, [-1, -1])]
    gen_task(ctx, targets, preds)


def gen_task7(upreds=None):
  """Logical OR: p(X):-q(X).p(X):-r(X)."""
  return logical_or(upreds=upreds)

def gen_task8(upreds=None):
  """Transitive case: p(X,Y):-q(X,Z);r(Z,Y)."""
  assert ARGS.arity == 2, "Transitive task arity must be 2."
  preds = r_preds(3, upreds)
  # Existential variable with single choice
  vs = r_vars(3)
  rule = [(preds[0], [vs[0], vs[2]]),
          (preds[1], vs[:2]),
          (preds[2], vs[1:])]
  if upreds:
    return rule
  ctx = [rule]
  # Add matching ground cases
  args = r_consts(3)
  add_pred(ctx, (preds[1], args[:2]), preds, args, 1.0)
  succ_supps = [0, 1, len(ctx), -1]
  add_pred(ctx, (preds[2], args[1:]), preds, args, 1.0)
  # Add non-matching ground cases
  argso = r_consts(3)
  argso.insert(R.randint(1, 2), r_consts(1, argso)[0])
  add_pred(ctx, (preds[1], argso[:2]), preds, argso, 0.5)
  add_pred(ctx, (preds[2], argso[2:]), preds, argso, 0.5)
  # Successful case
  # Fail on half-matching existential
  targets = [((preds[0], [args[0], args[2]]), 1, succ_supps),
             ((preds[0], [argso[0], argso[3]]), 0, [0, -1, -1, -1])]
  gen_task(ctx, targets, preds)

def gen_task9(upreds=None):
  """Single step deduction with NBF: p(X):--q(X)."""
  return nstep_deduction(1, True, upreds)

def gen_task10(upreds=None):
  """Double step deduction with NBF: p(X):--q(X).q(X):-r(X)."""
  return nstep_deduction(2, True, upreds)

def gen_task11(upreds=None):
  """Logical AND with NBF: p(X):-q(X);-r(X)."""
  return logical_and(True, upreds)

def gen_task12(upreds=None):
  """Logical OR with NBF: p(X):--q(X).p(X):-r(X)."""
  return logical_or(True, upreds)

def gen_task0():
  """Generate an ILP task example."""
  argc = 1
  goal = 'f'
  premise = 'b'
  ctx, targets = list(), list()
  # Generate according to goal <- premise
  args = r_consts(argc)
  # Add the successful ground case
  ctx.append([(premise, args)])
  targets.append(((goal, args), 1))
  # Fail on non-matching constant
  args = args.copy()
  args[R.randrange(len(args))] = r_consts(1, args)[0]
  preds = r_preds(3)
  ctx.append([(preds[0], args)])
  targets.append(((goal, args), 0))
  # Add padding length dummy rule
  vs = r_vars(argc)
  ctx.append([(preds[1], vs), (preds[2], vs)])
  preds.extend([goal, premise])
  gen_task(ctx, targets, preds)

if __name__ == '__main__':
  # pylint: disable=line-too-long
  # Arguments
  parser = argparse.ArgumentParser(description="Generate logic program data.")
  parser.add_argument("-t", "--task", default=1, type=int, help="The task to generate.")
  parser.add_argument("-s", "--size", default=1, type=int, help="Number of programs to generate.")
  # Configuration parameters
  parser.add_argument("-ns", "--noise_size", default=2, type=int, help="Size of added noise rules.")
  parser.add_argument("-cl", "--constant_length", default=1, type=int, help="Length of constants.")
  parser.add_argument("-vl", "--variable_length", default=1, type=int, help="Length of variables.")
  parser.add_argument("-pl", "--predicate_length", default=1, type=int, help="Length of predicates.")
  parser.add_argument("-ar", "--arity", default=2, type=int, help="Arity.")
  # Task specific options
  parser.add_argument("--nstep", type=int, help="Generate nstep deduction programs.")
  ARGS = parser.parse_args()

  # Generate given task
  task = "gen_task" + str(ARGS.task)
  for _ in range(ARGS.size):
    if ARGS.nstep:
      nstep_deduction(ARGS.nstep)
    else:
      globals()[task]()
