not really working black knot solver for 2006 icfp

knot_solve.py

#!/usr/bin/env python3.8

import sys
import re

def extract(s):
    return [int(x) for x in re.findall(r'-?\d+', s)]

from z3 import Bool, Int, If, Not, Solver, And

def solve(pieces):
    width = len(pieces)
    height = width*5 # XXX: sure??

    s = Solver()

    swaps = [[Bool(f's{r}_{c}') for c in range(width)] for r in range(height-1)]
    balls = [[Int(f'b{r}_{c}') for c in range(width)] for r in range(height)]
    plinks = [[Int(f'p{r}_{c}') for c in range(width)] for r in range(height)]

    for c in range(width):
        s.add(balls[0][c] == c)
        s.add(plinks[0][c] == 0)

    for r in range(height-1):
        for c in range(1, width):
            s.add(Not(And(swaps[r][c-1], swaps[r][c])))
        s.add(swaps[r][-1] == False)

    for r in range(1, height):
        for c in range(width):

            ball = balls[r-1][c]
            plink = plinks[r-1][c]
            if c > 0:

                lswap = swaps[r-1][c-1]
                ball = If(lswap, balls[r-1][c-1], ball)
                plink = If(lswap, plinks[r-1][c-1] + 1, plink)

            if c + 1 < width:
                rswap = swaps[r-1][c]
                ball = If(rswap, balls[r-1][c+1], ball)
                plink = If(rswap, plinks[r-1][c+1], plink)

            s.add(balls[r][c] == ball)
            s.add(plinks[r][c] == plink)

    for orig, target, plink in pieces:
        s.add(balls[-1][target] == orig)
        s.add(plinks[-1][target] == plink)

    print(s)
    s.check()

    model = s.model()
    print(model)

    return [[model[x] for x in r] for r in swaps]


def draw(model):
    return "\n".join([
        "".join([
            ">" if x else "<" if c > 0 and line[c-1] else "|"
            for c, x in enumerate(line)
        ])
        for line in model
    ])



def main(args):
    data = [extract(x.strip()) for x in sys.stdin]
    model = solve(data)
    print(model)
    print(draw(model))


if __name__ == '__main__':
    main(sys.argv)