#! /usr/bin/env python3

# Jaco van de Pol, Aarhus, 22 juli 2023

# Solve the "swapple" game as defined by Søren Fuglede (Kvantify)
# - input: a goal matrix (full rank)
# - goal: transform ID matrix to goal matrix by row/column additions (xor)
# - output: minimal number of additions

import time
import random
import sys
import z3
from z3_pysat import PysatSolver
from z3 import Bool, And, Or, Not, Xor, Implies
from math import log2, ceil

N = -1  # global variable for the dimension of the NxN matrix
logN = -1


def setN(n):
    global N, logN
    N = n
    logN = ceil(log2(N))


# Conventions:
# - M is any matrix (list of lists)
# - B is a matrix of Booleans (Python)
# - V is a matrix of Variables (Z3)
# - t is a time stamp
# - d is maximal depth
# - r/c is a row/column index


# transpose any square matrix
def transpose(M):
    return [[M[c][r] for c in range(N)] for r in range(N)]


# return identity matrix of booleans
def identity():
    return [[False] * i + [True] + [False] * (N - i - 1) for i in range(N)]


# print matrix of booleans as a 0/1 matrix
def print_matrix(B):
    for r in B:
        zero_ones = map(lambda x: "1" if x == True else ("0" if x == False else "?"), r)
        print(" ".join(zero_ones))


#
# DEFINE SIMPLE PATH CONSTRAINTS
################################

# uses existential bitvector variables to specify that matrices represented by V
# are different at time t0 and t1 at (row, col) = (DR[t0][t1], DC[t0][t1]).


def valN(x):  # return bitvector to represent constant x in logN bits
    return z3.BitVecVal(x, logN)


def simple_path_constraint(t, Vars, solver):
    for p in range(t):  # new state t different from all old states p < t
        DRpt = z3.BitVec(f"DR_{p}_{t+1}", logN)  # row where they differ
        DCpt = z3.BitVec(f"DC_{p}_{t+1}", logN)  # column where they differ
        solver.add(And(z3.ULE(DRpt, valN(N - 1)), z3.ULE(DCpt, valN(N - 1))))
        for r in range(N):
            for c in range(N):
                solver.add(
                    Implies(
                        And(DRpt == valN(r), DCpt == valN(c)),
                        Vars[p][r][c] != Vars[t + 1][r][c],
                    )
                )


# if two subsequent actions are independent (i.e., commute), they should be ordered
def partial_order_reduction(t, solver):
    # at t0: R0_2 := R0_2 + R0_1
    # at t1: R1_2 := R1_2 + R1_1
    #
    # independent: x := x+y, x := x+z
    # independent: x := x+y, z := z+w
    R0_1 = z3.BitVec(f"r1_{t-1}", logN)
    R0_2 = z3.BitVec(f"r2_{t-1}", logN)
    R1_1 = z3.BitVec(f"r1_{t}", logN)
    R1_2 = z3.BitVec(f"r2_{t}", logN)
    independent = Or(R0_2 == R1_2, And(R0_2 != R1_1, R1_2 != R0_1))
    ordered = And(z3.ULE(R0_1, R1_1), Implies(R0_1 == R1_1, z3.ULE(R0_2, R1_2)))
    solver.add(Implies(independent, ordered))


# Alternative 1: only check that the previous action is not "undone". Cheaper but less effective.
#    if t!=0: # avoid going back
#        solver.add(Or(R1!=z3.BitVec(f"r1_{t-1}", logN)),
#                      R2!=z3.BitVec(f"r2_{t-1}", logN))
#
# Alternative 2: compare the whole matrix, without using extra (bitvector) variables (seems slower).
#    for p in range(t-1): # new state t different from all old states p < t-1
#        Vp, Vt = vars(p), vars(t)
#        solver.add(Or([Vp[r][c]!=Vt[r][c] for c in range(N) for r in range(N)]))


# DEFINE THE ENCODING
#####################


# return matrix of variables at time t
def vars(t):
    return [[Bool(f"x_{r}_{c}_{t}") for c in range(N)] for r in range(N)]


# return a formula that defines V==B
def define_matrix(V, B):
    return And(
        [
            (
                V[r][c]
                if B[r][c] == True
                else (Not(V[r][c]) if B[r][c] == False else True)
            )
            for c in range(N)
            for r in range(N)
        ]
    )


# return a formula that defines V1 == V0[r2 := r2 + r1] (row-wise)
def add_row(V0, V1, r1, r2):
    return And(
        And([V1[r2][c] == Xor(V0[r1][c], V0[r2][c]) for c in range(N)]),
        And([V1[r][c] == V0[r][c] for r in range(N) if r != r2 for c in range(N)]),
    )


# return a formula that adds some row to some other row
def add_some_row(V0, V1):
    return Or(
        [add_row(V0, V1, r1, r2) for r1 in range(N) for r2 in range(N) if r1 != r2]
    )


# add bounded-model-check constraints to the Solver to go from init to goal in d steps
# version: Booleans, nested and-or-and
def BMC_bool(depth, init, goal, cycle, solver):
    solver.add(define_matrix(vars(0), init))
    solver.add(define_matrix(vars(depth), goal))
    Vars = [vars(t) for t in range(depth + 1)]
    for t in range(depth):
        solver.add(add_some_row(Vars[t], Vars[t + 1]))
        if cycle:
            simple_path_constraint(t, Vars, solver)


### Bitvectors


# add bounded-model-check constraints to the Solver to go from init to goal in d steps
# version: Bit-vectors
def BMC_bitvector(depth, init, goal, cycle, por, solver):
    Vars = [vars(t) for t in range(depth + 1)]
    solver.add(define_matrix(Vars[0], init))
    solver.add(define_matrix(Vars[depth], goal))
    for t in range(depth):
        R0 = z3.BitVec(f"r1_{t}", logN)  # row to add from
        R1 = z3.BitVec(f"r2_{t}", logN)  # row to add to
        V0, V1 = Vars[t], Vars[t + 1]
        solver.add(
            And(z3.ULE(R0, valN(N - 1)), z3.ULE(R1, valN(N - 1)), R0 != R1)
        )  # ULE = Unsigned Less-or-Equal
        for r2 in range(N):
            for c in range(N):
                for r1 in range(N):
                    solver.add(
                        Implies(
                            And(R1 == valN(r2), R0 == valN(r1)),
                            V1[r2][c] == Xor(V0[r1][c], V0[r2][c]),
                        )
                    )
                solver.add(Implies(R1 != valN(r2), V1[r2][c] == V0[r2][c]))
        if cycle:
            simple_path_constraint(t, Vars, solver)
        if t > 0 and por:
            partial_order_reduction(t, solver)


# add bounded-model-check constraints to the Solver to go from init to goal in d steps
# version: Booleans, use one-hot encoding to choose which row
def BMC_onehot(depth, init, goal, cycle, solver):
    solver.add(define_matrix(vars(0), init))
    solver.add(define_matrix(vars(depth), goal))
    Vars = [vars(t) for t in range(depth + 1)]
    for t in range(depth):
        V0, V1 = Vars[t], Vars[t + 1]
        R1 = [z3.Bool(f"r1_{r}_{t}") for r in range(N)]  # row to add from
        R2 = [z3.Bool(f"r2_{r}_{t}") for r in range(N)]  # row to add to
        solver.add(z3.AtMost(R1 + [1]))
        solver.add(z3.AtMost(R2 + [1]))
        solver.add(z3.AtLeast(R1 + [1]))
        solver.add(z3.AtLeast(R2 + [1]))
        for r2 in range(N):
            solver.add(Not(And(R1[r2], R2[r2])))
            for c in range(N):
                for r1 in range(N):
                    solver.add(
                        Implies(
                            And([R2[r2], R1[r1]]),
                            V1[r2][c] == Xor(V0[r1][c], V0[r2][c]),
                        )
                    )
                solver.add(Implies(And(Not(R2[r2])), V1[r2][c] == V0[r2][c]))
        if cycle:
            simple_path_constraint(t, Vars, solver)


# Find the shortest path of maximal depth to the goal matrix in Z3 solver
# returns -1 if no path is found
def solve_swapple(goal, start, depth, cycle, por, encoding, solver):
    init = identity()
    for t in range(start, depth + 1):
        print(f"try t={t}", end=" ", flush=True)
        start = time.perf_counter()
        solver.reset()  # maybe could do incremental solving?
        if encoding == "bool":
            BMC_bool(t, init, goal, cycle, solver)
        elif encoding == "bitvec":
            BMC_bitvector(t, init, goal, cycle, por, solver)
        elif encoding == "onehot":
            BMC_onehot(t, init, goal, cycle, solver)
        else:
            assert False, "Error (internal): encoding not understood"
        # for x in solver.assertions():
        #    print(z3.simplify(x))
        mid = time.perf_counter()
        print(f"({mid-start:.2f}s", end=" ", flush=True)
        found = solver.check()
        end = time.perf_counter()
        print(f"+ {end-mid:.2f}s)", end=" ")
        if found == z3.sat:
            print(f"yes!")
            return t
        elif found == z3.unsat:
            print(f"no")
        else:  # for instance after Ctrl-C
            print(f"don't know")
    return -1


# CHECK AND RETRIEVE THE RESULT
###############################


# Return matrix with row r1 added to row r2
def xor_matrix(B, r1, r2):
    return B[:r2] + [[B[r1][c] ^ B[r2][c] for c in range(N)]] + B[r2 + 1 : N]


# compute some action that transforms B0 to B1
def retrieve_step(B0, B1):
    for r1 in range(N):
        for r2 in range(N):
            if B1 == xor_matrix(B0, r1, r2):
                return ("row", r1, r2)
    print_matrix(B0)
    print("->")
    print_matrix(B1)
    assert False, "Error (internal): step not recognized!"


# Print the actions (and possibly the matrices) from the model
def print_steps(model, t, verbose=False):
    matrices = []
    for k in range(t + 1):
        V = vars(k)
        B = [[z3.is_true(model[V[r][c]]) for c in range(N)] for r in range(N)]
        matrices.append(B)
    if verbose:
        print_matrix(matrices[0])
    for k in range(1, t + 1):
        (str, x, y) = retrieve_step(matrices[k - 1], matrices[k])
        print(f"Step {k}: add {str} {x+1} to {y+1}")
        if verbose:
            print_matrix(matrices[k])


# Generate goal matrices automatically
######################################


# enumerate all subsets of L (as list)
def subsets(L):
    if len(L) == 0:
        yield []
    else:
        for x in subsets(L[1:]):
            yield x
            yield [L[0]] + x


# add all rows in B by row-wise xor
def xorall(L):
    result = [False] * N
    for row in L:
        for i in range(N):
            result[i] ^= row[i]
    return result


# Check if row is independent from all other rows
def independent(row, rows):
    for L in subsets(rows):
        if xorall(L) == row:
            return False
    return True


def filter_indep(rows, iterator):
    for row in iterator:
        if independent(row, rows):
            yield row


# Enumerate all rows of length
def gen_row(n):
    if n == 0:
        yield []
    else:
        for row in gen_row(n - 1):
            for a in (False, True):
                yield [a] + row


# Enumerate all m x N matrices of full rank
def gen_all(m):
    if m == 0:
        yield []
    else:
        for M in gen_all(m - 1):
            for row in filter_indep(M, gen_row(N)):
                yield M + [row]


# Enumerate all N x N full-rank matrices
def enumerate_all():
    print(f"Generating all matrices of size {N} x {N}")
    yield from gen_all(N)


# Enumerate num random NxN matrices of full rank by "shuffling" random actions
def random_swapple(seed, shuffle, num, verbose=False):
    random.seed(seed)
    for i in range(num):
        B = identity()
        print(f"Create a matrix of size {N}x{N} ({shuffle} shuffles).")
        for t in range(shuffle):
            if verbose:
                print_matrix(B)
            rc = random.randrange(2)  # add row or column
            r1 = random.randrange(N)
            r2 = random.randrange(N - 1)  # r2 must be different from r1
            if r2 >= r1:
                r2 += 1
            print(f"Shuffle {t+1}: add {'row' if rc==0 else 'col'} {r1+1} to {r2+1}")
            if rc == 0:
                B = xor_matrix(B, r1, r2)
            else:
                B = transpose(xor_matrix(transpose(B), r1, r2))
        yield B


# enumerate all permutations of 1..N
def gen_perm(L):
    if len(L) == 0:
        yield []
    else:
        for i in range(len(L)):
            for pi in gen_perm(L[:i] + L[i + 1 :]):
                yield [L[i]] + pi


def enum_perm():
    yield from gen_perm(list(range(N)))


def enum_all_perm():
    print(f"Generating all matrices of size {N} x {N}")
    for pi in enum_perm():
        B = []
        for i in pi:
            B.append([False] * i + [True] + [False] * (N - i - 1))
        yield B


# Main Functions
################


# Read a square array of 0/1 from file (space separated)
# Return a matrix, i.e. a list of list of Booleans (or None for don't care)
def read_matrix(file, name):
    B = []
    for line in file:
        if line[0] not in "%\n":  # skip comments and empty lines
            row = list(
                map(
                    lambda s: (
                        True if s[0] == "1" else (False if s[0] == "0" else None)
                    ),
                    line.split(" "),
                )
            )
            B.append(row)
            if N == -1:
                setN(len(row))
            else:
                assert N == len(row), "Error (input): Rows should have equal length"
    assert N == len(B), f"Error (input): Should be a square matrix {N}!={len(B)}"
    print(f"Read a matrix of size {N}x{N} from {name}.")
    return B


# Solve swapple goal with solver, and print the result
def swapple_main(goal, solver, args):
    if args.depth == 0:
        args.depth = 3 * (N - 1)
    print(f"\nUsing {args.encoding} encoding,", end=" ")
    print(f"with{'' if args.cycle else 'out'} simple path constraint", end=" ")
    if args.encoding == "bitvec":
        print(f"and with{'' if args.por else 'out'} POR")
    print(f"\nUsing SAT solver {args.satsolver} for depth {args.start}-{args.depth}")
    shortest = solve_swapple(
        goal, args.start, args.depth, args.cycle, args.por, args.encoding, solver
    )
    if shortest >= 0:
        print_steps(solver.model(), shortest, args.verbose)
        print(f"\nSolved {N} x {N} matrix in {shortest} steps")
        return shortest
    else:
        print(f"\nNo solution found after {args.depth} steps")
        return args.depth + 1


def swapple_experiment(goals, solver, args):
    def text(i):
        print(f"\n==============\nExperiment {i}:\n==============")

    solved = dict()
    i = 1
    text(i)
    for goal in goals:
        print_matrix(goal)
        d = swapple_main(goal, solver, args)
        if d in solved:
            solved[d] += 1
        else:
            solved[d] = 1
        i += 1
        text(i)
    print("Finished all goals")
    return solved


def print_freq_table(solved, total=None):
    exp = 0
    for i in sorted(solved):
        if i <= args.depth:
            exp += solved[i]
            print(f"solved in {i} steps: {solved[i]}")
    d = args.depth + 1
    if d in solved:
        exp += solved[d]
        print(f"unsolved after {d-1} steps: {solved[d]}")
    assert (
        total == None or total == exp
    ), f"Error (internal): experiment counts don't match ({total}, {exp})"
    print(f"({exp} experiments in total)")


def extendStrict(m):
    result = []
    for row in m:
        result.append(row + [False])
    extra = [False] * len(m) + [True]
    result.append(extra)
    return result


def extendMedium(m):
    result = []
    for row in m:
        result.append(row + [False])
    extra = [None] * (len(m) + 1)
    result.append(extra)
    return result


def extendRelaxed(m):
    result = []
    for row in m:
        result.append(row + [None])
    extra = [None] * (len(m) + 1)
    result.append(extra)
    return result


def test_conjecture(solver, args):
    print(f"Testing conjecture with N={N}")
    count = 0
    for m in enumerate_all():
        count += 1
        print_matrix(m)
        n = swapple_main(m, solver, args)
        print(f"Conjecture {N}x{N}: solved in {n} steps ({count})")
        m1 = extendMedium(m)
        oldN = N
        setN(N + 1)
        n1 = swapple_main(m1, solver, args)
        print(f"Conjecture {N}x{N}: solved in {n1} steps ({count})")
        if n1 != n:
            print("Conjecture: Counter example found!")
            print_matrix(m1)
            exit()
        setN(oldN)


if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser(
        description="Solve a SWAPPLE instance with a SAT solver"
    )
    parser.add_argument(
        "-v", "--verbose", help="show intermediate matrices", action="store_true"
    )
    parser.add_argument(
        "-e",
        "--encoding",
        help="Used BMC encoding (bitvec)",
        choices=["bool", "bitvec", "onehot"],
        default="bitvec",
    )
    parser.add_argument(
        "-c",
        "--cycle",
        help="Avoid cycles by adding simple-path-constraint (default 1)",
        type=int,
        default=1,
    )
    parser.add_argument(
        "--por", help="Partial-Order Reduction (default 1)", type=int, default=1
    )

    mode = parser.add_mutually_exclusive_group()
    mode.add_argument(
        "-r",
        "--random",
        help="generate random matrix of size N x N",
        metavar="N",
        type=int,
        default=0,
    )
    mode.add_argument(
        "-a",
        "--all",
        help="treat all full-rank matrices of size N x N",
        metavar="N",
        type=int,
        default=0,
    )
    mode.add_argument(
        "-p",
        "--perm",
        help="treat all permutation matrices of size N x N",
        metavar="N",
        type=int,
        default=0,
    )
    mode.add_argument("-f", "--file", help="read input matrix from file")
    mode.add_argument(
        "--conjecture",
        help="(tmp) test conjecture: ancillary not needed",
        metavar="N",
        type=int,
        default=0,
    )

    rnd = parser.add_argument_group("for random mode")
    rnd.add_argument(
        "-s", "--seed", help="seed for random generator (None)", default=None
    )
    rnd.add_argument(
        "-sh",
        "--shuffle",
        help="number of shuffles for random matrix (15)",
        metavar="SH",
        default=15,
        type=int,
    )
    rnd.add_argument("-l", "--loop", help="repeat experiment (1)", default=1, type=int)
    solver = parser.add_argument_group("for solving")
    solver.add_argument(
        "-d",
        "--depth",
        help="Maximum depth for bounded model check. default: 3(N-1)",
        metavar="D",
        type=int,
        default=0,
    )
    solver.add_argument(
        "--start", help="Minimum depth considered (default 0)", default=0, type=int
    )
    solver.add_argument(
        "--satsolver",
        help="z3, z3blast or any sat solver from PySat (z3blast)",
        metavar="SOLVER",
        default="z3blast",
    )

    args = parser.parse_args()

    if args.verbose:
        print(args)

    # Determine the solver
    if args.satsolver == "z3":
        Solver = z3.Solver()
    elif args.satsolver == "z3blast":
        Solver = z3.Then("simplify", "bit-blast", "sat").solver()
    else:
        Solver = PysatSolver(args.satsolver)

    if args.file != None:
        mode = "file"
    elif args.random > 0:
        mode = "random"
    elif args.all > 0:
        mode = "all"
    elif args.perm > 0:
        mode = "perm"
    elif args.conjecture > 0:
        mode = "conjecture"
    else:
        mode = "stdin"

    if mode == "file":
        try:
            goal = read_matrix(open(args.file), args.file)  # this also defines N
            swapple_main(goal, Solver, args)
        except FileNotFoundError:
            print(f"Error (input): File {args.file} not found")
    elif mode == "stdin":
        print("reading matrix from stdin:")
        goal = read_matrix(sys.stdin, "stdin")  # this also defines N
        swapple_main(goal, Solver, args)
    elif mode == "random":
        setN(args.random)
        goals = random_swapple(args.seed, args.shuffle, args.loop, args.verbose)
        solved = swapple_experiment(goals, Solver, args)
        print(
            f"\nFrequency Table ({N} x {N} matrix, {args.shuffle} shuffles, seed={args.seed}, depth {args.depth})"
        )
        print_freq_table(solved, args.loop)
    elif mode == "all":
        setN(args.all)
        goals = enumerate_all()
        solved = swapple_experiment(goals, Solver, args)
        print(
            f"\nFrequency Table (all {N} x {N} full-rank matrices, depth {args.depth})"
        )
        print_freq_table(solved)
    elif mode == "perm":
        setN(args.perm)
        goals = enum_all_perm()
        solved = swapple_experiment(goals, Solver, args)
        print(
            f"\nFrequency Table (all {N} x {N} permutation matrices, depth {args.depth})"
        )
        print_freq_table(solved)
    elif mode == "conjecture":
        setN(args.conjecture)
        test_conjecture(Solver, args)
    else:
        print(f"Error: nothing to do for mode {mode}")
