NumbersLink_CSP.py

'''
"Numbers Link"パズルを制約充足問題(CSP)ソルバーSugarで解く・作る
'''
import argparse
import os, sys
import random
import time
import re
import subprocess


def parse_map(data):
    # サイズを取得
    width = max([len(x) for x in data])
    height = len(data)
    # 数字をチェック
    nums = {}
    num_map = [[0 for x in range(width)] for y in range(height)]
    for iy, line in enumerate(data):
        for ix, val in enumerate(line):
            ival = int(val)
            if ival not in nums:
                nums[ival] = []
            nums[ival].append((ix, iy))
            num_map[iy][ix] = ival
    return width, height, nums, num_map


def create_map(width, height, max_num, size):
    num_map = [[0 for x in range(width)] for y in range(height)]
    for num in range(1, max_num+1):
        cnt = 0
        while cnt < size:
            ix = random.randint(0, width-1)
            iy = random.randint(0, height-1)
            if num_map[iy][ix] == 0:
                num_map[iy][ix] = num
                cnt += 1
    return num_map


def save_csp(file, width, height, nums, num_map):
    with open(file, mode='w') as f:
        # 横方向の辺h
        for iy in range(height):
            for ix in range(width-1):
                f.write('(int hr_{}_{} 0 1)\n'.format(ix, iy))
                f.write('(int hl_{}_{} 0 1)\n'.format(ix, iy))
                f.write('(<= (+ hr_{0}_{1} hl_{0}_{1}) 1)\n'.format(ix, iy))
        # 縦方向の辺v
        for iy in range(height-1):
            for ix in range(width):
                f.write('(int vu_{}_{} 0 1)\n'.format(ix, iy))
                f.write('(int vd_{}_{} 0 1)\n'.format(ix, iy))
                f.write('(<= (+ vu_{0}_{1} vd_{0}_{1}) 1)\n'.format(ix, iy))

        # 節点
        for iy in range(height):
            for ix in range(width):
                # 節点nの入次数
                f.write('(int di_{}_{} 0 1)\n'.format(ix, iy))
                # 節点nの出次数
                f.write('(int do_{}_{} 0 1)\n'.format(ix, iy))
                # 節点nの入次数は周辺から入る辺の和
                hvs = []
                if ix > 0:
                    hvs.append('hr_{}_{}'.format(ix-1, iy))
                if ix < width-1:
                    hvs.append('hl_{}_{}'.format(ix, iy))
                if iy > 0:
                    hvs.append('vd_{}_{}'.format(ix, iy-1))
                if iy < height-1:
                    hvs.append('vu_{}_{}'.format(ix, iy))
                f.write('(= di_{}_{} (+ {}))\n'.format(
                    ix, iy, ' '.join(hvs)
                ))
                # 節点nの出次数は周辺へ出る辺の和
                hvs = []
                if ix > 0:
                    hvs.append('hl_{}_{}'.format(ix-1, iy))
                if ix < width-1:
                    hvs.append('hr_{}_{}'.format(ix, iy))
                if iy > 0:
                    hvs.append('vu_{}_{}'.format(ix, iy-1))
                if iy < height-1:
                    hvs.append('vd_{}_{}'.format(ix, iy))
                f.write('(= do_{}_{} (+ {}))\n'.format(
                    ix, iy, ' '.join(hvs)
                ))
                # 節点nの順番
                f.write('(int o_{}_{} -1 {})\n'.format(
                    ix, iy, width * height
                ))

        # 節点nの数字を表すドメインnumber
        f.write('(domain number ({}))\n'.format(
            ' '.join([str(n) for n in nums.keys()])
        ))
        # 節点n
        for num in nums:
            # 節点nに数字がない場合
            if num == 0:
                for pos in nums[num]:
                    # 節点nの入次数と出次数は等しい
                    f.write('(= di_{0}_{1} do_{0}_{1})\n'.format(pos[0], pos[1]))
                    # 節点nの数字はnumbers
                    f.write('(int x_{}_{} number)\n'.format(pos[0], pos[1]))
                    # 節点nの初期数字は0
                    f.write('(int i_{}_{} 0)\n'.format(pos[0], pos[1]))
                    # 節点nの数字が0の場合、入次数と出次数は0
                    f.write('(=> (= x_{0}_{1} 0) (= di_{0}_{1} 0))\n'.format(
                        pos[0], pos[1]
                    ))
                    f.write('(=> (= x_{0}_{1} 0) (= do_{0}_{1} 0))\n'.format(
                        pos[0], pos[1]
                    ))
                    # 節点nの数字が0の場合、順序は-1
                    f.write('(=> (= x_{0}_{1} 0) (= o_{0}_{1} -1))\n'.format(
                        pos[0], pos[1]
                    ))

            # 節点nに数字num(>0)がある場合
            else:
                dis = []
                dos = []
                dss  = []
                for pos in nums[num]:
                    # 節点nの数字はnum
                    f.write('(int x_{}_{} {})\n'.format(
                        pos[0], pos[1], num
                    ))
                    # 節点nの初期数字はnum
                    f.write('(int i_{}_{} {})\n'.format(
                        pos[0], pos[1], num
                    ))
                    # 節点の次数をリストに追加
                    dis.append('di_{}_{}'.format(pos[0], pos[1]))
                    dos.append('do_{}_{}'.format(pos[0], pos[1]))
                    dss.append('(+ di_{0}_{1} do_{0}_{1})'.format(pos[0], pos[1]))
                    # 始点の順番は0
                    f.write((
                        '(=> (and (= di_{0}_{1} 0) (= do_{0}_{1} 1)) '
                        '(= o_{0}_{1} 0))\n'
                    ).format(pos[0], pos[1]))
                    # 終点だった場合は順番はlen(nums[num])-1
                    f.write((
                        '(=> (and (= di_{0}_{1} 1) (= do_{0}_{1} 0)) '
                        '(= o_{0}_{1} {2}))\n'
                    ).format(pos[0], pos[1], len(nums[num])-1))
                # 入次数のうち1つが0でその他は1、かつ、
                # 出次数のうち1つが0でその他が1、かつ、
                # 入次数と出自数の和のうち2つが1でその他が2
                f.write('(count 0 ({}) eq 1)\n'.format(' '.join(dis)))
                f.write('(count 0 ({}) eq 1)\n'.format(' '.join(dos)))
                f.write('(count 1 ({}) eq 2)\n'.format(' '.join(dss)))

        # 横方向の辺h
        for iy in range(height):
            for ix in range(width-1):
                # 辺で接続された節点の数字は等しい
                f.write('(=> (> hr_{0}_{1} 0) (= x_{0}_{1} x_{2}_{1}))\n'.format(
                    ix, iy, ix+1
                ))
                f.write('(=> (> hl_{0}_{1} 0) (= x_{0}_{1} x_{2}_{1}))\n'.format(
                    ix, iy, ix+1
                ))
                # 辺で接続された節点の順序
                # 接続先の初期数が0の場合、順序は等しい
                f.write((
                    '(=> (and (> hr_{0}_{1} 0) (= i_{2}_{1} 0)) '
                    '(= o_{0}_{1} o_{2}_{1}))\n'
                ).format(ix, iy, ix+1))
                f.write((
                    '(=> (and (> hl_{0}_{1} 0) (= i_{0}_{1} 0)) '
                    '(= o_{0}_{1} o_{2}_{1}))\n'
                ).format(ix, iy, ix+1))
                # 接続先の初期数が>0の場合、順序は1増える
                f.write((
                    '(=> (and (> hr_{0}_{1} 0) (> i_{2}_{1} 0)) '
                    '(= o_{2}_{1} (+ o_{0}_{1} 1)))\n'
                ).format(ix, iy, ix+1))
                f.write((
                    '(=> (and (> hl_{0}_{1} 0) (> i_{0}_{1} 0)) '
                    '(= o_{0}_{1} (+ o_{2}_{1} 1)))\n'
                ).format(ix, iy, ix+1))
        # 縦方向の辺v
        for iy in range(height-1):
            for ix in range(width):
                # 辺で接続された節点の数字は等しい
                f.write('(=> (> vu_{0}_{1} 0) (= x_{0}_{1} x_{0}_{2}))\n'.format(
                    ix, iy, iy+1
                ))
                f.write('(=> (> vd_{0}_{1} 0) (= x_{0}_{1} x_{0}_{2}))\n'.format(
                    ix, iy, iy+1
                ))
                # 辺で接続された節点の順を更新
                # 接続先の初期数が0の場合
                f.write((
                    '(=> (and (> vu_{0}_{1} 0) (= i_{0}_{1} 0)) '
                    '(= o_{0}_{1} o_{0}_{2}))\n'
                ).format(ix, iy, iy+1))
                f.write((
                    '(=> (and (> vd_{0}_{1} 0) (= i_{0}_{2} 0)) '
                    '(= o_{0}_{1} o_{0}_{2}))\n'
                ).format(ix, iy, iy+1))
                # 接続先の初期数が>0の場合
                f.write((
                    '(=> (and (> vu_{0}_{1} 0) (> i_{0}_{1} 0)) '
                    '(= o_{0}_{1} (+ o_{0}_{2} 1)))\n'
                ).format(ix, iy, iy+1))
                f.write((
                    '(=> (and (> vd_{0}_{1} 0) (> i_{0}_{2} 0)) '
                    '(= o_{0}_{2} (+ o_{0}_{1} 1)))\n'
                ).format(ix, iy, iy+1))


def call_cmd(cmd):
    return subprocess.check_output(
        cmd,
        shell=True
    ).decode('utf-8').splitlines()


def append_csp(file, result):
    with open(file, mode='a') as f:
        res_h = [[0 for ix in range(width-1)] for iy in range(height)]
        res_v = [[0 for ix in range(width)] for iy in range(height-1)]
        for line in result:
            m = re.match(r'a (hr|hl)_([0-9])+_([0-9]+)\t([0-9]+)', line)
            if m is not None and m.group(4) == '1':
                res_h[int(m.group(3))][int(m.group(2))] = 1
            m = re.match(r'a (vu|vd)_([0-9])+_([0-9]+)\t([0-9]+)', line)
            if m is not None and m.group(4) == '1':
                res_v[int(m.group(3))][int(m.group(2))] = 1
        cond = []
        for iy, line in enumerate(res_h):
            for ix, val in enumerate(line):
                if val == 0:
                    cond.append('(and (= hr_{0}_{1} 0) (= hl_{0}_{1} 0))'.format(
                        ix, iy
                    ))
                else:
                    cond.append(
                        (
                            '(or (and (= hr_{0}_{1} 0) (= hl_{0}_{1} 1)) '
                            '(and (= hr_{0}_{1} 1) (= hl_{0}_{1} 0)))'
                        ).format(
                            ix, iy
                        )
                    )
        for iy, line in enumerate(res_v):
            for ix, val in enumerate(line):
                if val == 0:
                    cond.append('(and (= vu_{0}_{1} 0) (= vd_{0}_{1} 0))'.format(
                        ix, iy
                    ))
                else:
                    cond.append(
                        (
                            '(or (and (= vu_{0}_{1} 0) (= vd_{0}_{1} 1)) '
                            '(and (= vu_{0}_{1} 1) (= vd_{0}_{1} 0)))'
                        ).format(
                            ix, iy
                        )
                    )
        f.write('(not (and {}))\n'.format(' '.join(cond)))


def to_lines(width, height, nums, num_map, result):
    lines = []
    res_n = [[' ' for ix in range(width)] for iy in range(height)]
    res_h = [[' ' for ix in range(width-1)] for iy in range(height)]
    res_v = [[' ' for ix in range(width)] for iy in range(height-1)]

    for line in result:
        # 節点
        m = re.match(r'a x_([0-9]+)_([0-9]+)\t([0-9]+)', line)
        if m is not None:
            ix = int(m.group(1))
            iy = int(m.group(2))
            num = int(m.group(3))
            if num > 0:
                if num_map[iy][ix] > 0:
                    res_n[iy][ix] = str(num)
                else:
                    res_n[iy][ix] = '+'
        # 横方向の辺
        m = re.match(r'a h[rl]_([0-9]+)_([0-9]+)\t([0-9]+)', line)
        if m is not None and m.group(3) == '1':
            res_h[int(m.group(2))][int(m.group(1))] = '-'
        # 縦方向の辺
        m = re.match(r'a v[ud]_([0-9]+)_([0-9]+)\t([0-9]+)', line)
        if m is not None and m.group(3) == '1':
            res_v[int(m.group(2))][int(m.group(1))] = '|'

    out = []
    for iy in range(height-1):
        out.clear()
        for ix in range(width-1):
            out.append(res_n[iy][ix])
            out.append(res_h[iy][ix])
        out.append(res_n[iy][width-1])
        lines.append(' '.join(out))

        out.clear()
        for ix in range(width):
            out.append(res_v[iy][ix])
        lines.append('   '.join(out))

    out.clear()
    for ix in range(width-1):
        out.append(res_n[height-1][ix])
        out.append(res_h[height-1][ix])
    out.append(res_n[height-1][height-1])
    lines.append(' '.join(out))

    return lines


def find_results(path_csp, max_loop=0):
    cnt = 0
    while max_loop == 0 or cnt < max_loop:
        # sugarを実行
        result = call_cmd('sugar {}'.format(path_csp))
        # 探索結果がある場合
        if result[0] == 's SATISFIABLE':
            cnt += 1
            # 表示用の経路文字列を作成
            lines = to_lines(width, height, nums, num_map, result)
            # 結果を出力
            yield lines
            # 結果の経路をCSPに追加し、別経路を探索
            append_csp(path_csp, result)
        # 探索結果がない場合
        else:
            return None



if __name__ == '__main__':

    parser = argparse.ArgumentParser(
        description='"Numbers Link" Solver/Creator with a CSP solver "sugar"',
        add_help=True
    )
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--solver', action = 'store_true', help = 'solver mode')
    group.add_argument('--creator', action = 'store_true', help = 'creator mode')
    # arguments for solver
    parser.add_argument('--map', required='--solver' in sys.argv, help='map data path')
    # arguments for creator
    parser.add_argument('--width', type=int, required='--creator' in sys.argv, help='map width')
    parser.add_argument('--height', type=int, required='--creator' in sys.argv, help='map height')
    parser.add_argument('--max_num', type=int, required='--creator' in sys.argv, help='maximum value of numbers')
    parser.add_argument('--size', type=int, required='--creator' in sys.argv, help='number of each numbers')
    parser.add_argument('--max_solve', type=int, required='--creator' in sys.argv, help='maximum number to solve')
    parser.add_argument('--num_loop', type=int, required='--creator' in sys.argv, help='number of challenge')
    parser.add_argument('--out', required='--creator' in sys.argv, help='output directory path')
    # common arguments
    parser.add_argument('--tmp', help='working directory path', required=True)
    args = parser.parse_args()

    # 問題を解く
    if args.solver:
        with open(args.map) as f:
            # 問題を読み込み
            data = [s.strip() for s in f.readlines()]
            width, height, nums, num_map = parse_map(data)
            print('width : {}'.format(width))
            print('height: {}'.format(height))
            print(''.join(['=' for x in range(4 * width - 3)]))
            # 時間計測開始
            start = time.time()
            # sugarに入力するためのCSPを出力
            path_csp = os.path.join(
                args.tmp,
                'solver.tmp{:05}.csp'.format(random.randint(10000, 99999))
            )
            save_csp(path_csp, width, height, nums, num_map)
            # 経路が探索できる限り、sugarを繰り返し実行
            cnt = 0
            for result in find_results(path_csp):
                if result is not None:
                    cnt += 1
                    # 結果を表示
                    print('[result: {}]'.format(cnt))
                    for line in result:
                        print(line)
                    print(''.join(['=' for x in range(4 * width - 3)]))
            # 時間計測終了
            elapsed_time = time.time() - start
            # 結果表示
            print('num satisfiable results = {}'.format(cnt))
            print('elapsed_time = {0:.3f}[sec]'.format(elapsed_time))

    # 問題を作る
    elif args.creator:
        if os.path.isdir(args.out):
            print('error: output directory already exists!')
        else:
            os.makedirs(args.out)
            path_csp = os.path.join(
                args.tmp,
                'creator.tmp{:05}.csp'.format(random.randint(10000, 99999))
            )
            for i in range(args.num_loop):
                # 問題をランダムに作成
                data = create_map(args.width, args.height, args.max_num, args.size)
                width, height, nums, num_map = parse_map(data)
                # sugarに入力するためのCSPを出力
                save_csp(path_csp, width, height, nums, num_map)
                # 経路が探索できる限り、sugarを繰り返し実行
                results = []
                for result in find_results(path_csp, args.max_solve):
                    if result is not None:
                        results.append(result)
                if len(results) > 0:
                    # 問題を出力
                    with open(
                        os.path.join(args.out, '{}_{:05}.txt'.format(len(results), i)),
                        mode='w'
                    ) as f:
                        for row in data:
                            f.write(''.join([str(x) for x in row]) + '\n')
                    # 結果を出力
                    with open(
                        os.path.join(args.out, '{}_{:05}_out.txt'.format(len(results), i)),
                        mode='w'
                    ) as f:
                        for result in results:
                            f.write(''.join(['=' for x in range(4 * width - 3)]) + '\n')
                            for line in result:
                                f.write(line + '\n')
                        f.write(''.join(['=' for x in range(4 * width - 3)]) + '\n')
                # 結果を表示
                print('{:05}: num satisfiable results = {}'.format(i, len(results)))