#! /usr/bin/env python3

# Irfansha Shaik, Aarhus, 31 August 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 additions (xor)
# - output: minimal number of additions
import sys
import swapple_z3 as sz
import os
import time


# Print the actions from the plan
def print_steps(plan, t, verbose=False):
    for k in range(1, t + 1):
        cur_action = plan[k - 1]
        print(f"Step {k}: add row {cur_action[1][1:]} to {cur_action[2][1:]}")


def parse_fdplan(plan_file):
    try:
        f = open(plan_file, "r")
        lines = f.readlines()
        f.close()
    except FileNotFoundError:
        print(f"No plan could be found. Error!")
        exit(-1)
    plan = []
    for line in lines:
        # only if not a commit:
        if ";" not in line:
            plan.append(line.strip(")\n").strip("()").split(" "))
            # print(line)
    return plan


def run_fdownward(args):
    # removing existing plan for correctness:
    if os.path.exists(args.plan_file):
        os.remove(args.plan_file)
    if args.planner == "fdss-sat":
        planner_options = "--alias seq-sat-fdss-2018 --portfolio-single-plan"
    elif args.planner == "fdss-opt":
        planner_options = "--alias seq-opt-fdss-1"
    elif args.planner == "fdss-opt-2":
        planner_options = "--alias seq-opt-fdss-2"
    elif args.planner == "fd-lmcut":
        planner_options = "--alias seq-opt-lmcut"
    elif args.planner == "fd-bjolp":
        planner_options = "--alias seq-opt-bjolp"
    elif args.planner == "fd-ms":
        planner_options = "--alias seq-opt-merge-and-shrink"

    if os.system("fast-downward.py -v >" + os.devnull) != 0:
        print(
            f"Error: planner {args.planner} requires executable 'fast-downward.py' on the path"
        )
        exit(-1)
    command = (
        "fast-downward.py "
        + planner_options
        + " --log-level warning --plan-file "
        + args.plan_file
        + " --sas-file "
        + args.SAS_file
        + "  --overall-time-limit "
        + str(int(args.time))
        + "s "
        + args.pddl_domain_out
        + " "
        + args.pddl_problem_out
        + " > "
        + args.log_out
    )
    if args.verbose > -1:
        print(command, flush=True)
    os.system(command)


def generate_problem_specification(goal):
    pddl_lines = []
    pddl_lines.append("(define (problem instance)\n")
    pddl_lines.append("  (:domain Swapple)\n")
    pddl_lines.append("  (:objects")
    # generating objects for rows/columns:
    objects = ""
    for i in range(1, sz.N + 1):
        objects += f"o{i} "
    pddl_lines.append(f"  {objects}- object)\n")

    pddl_lines.append("(:init\n")
    # identity matrix:
    cells = ""
    for i in range(1, sz.N + 1):
        cells += f"(cell o{i} o{i})\n"

    pddl_lines.append(cells)
    pddl_lines.append(")\n")

    pddl_lines.append("(:goal\n")

    pddl_lines.append("  (and ;; target matrix\n")
    for row_id in range(1, sz.N + 1):
        row_line = ""
        for column_id in range(1, sz.N + 1):
            cell = goal[row_id - 1][column_id - 1]
            if cell == True:
                row_line += f"     (cell o{row_id} o{column_id}) "
            elif cell == False:
                row_line += f"(not (cell o{row_id} o{column_id}))"
            else:
                assert cell == None
        pddl_lines.append(f" {row_line}\n")
    pddl_lines.append("  )\n)\n)")
    return pddl_lines


def print_freq_table(solved, total=None):
    exp = 0
    for i in sorted(solved):
        exp += solved[i]
        print(f"solved in {i} steps: {solved[i]}")
    assert (
        total == None or total == exp
    ), f"Error (internal): experiment counts don't match ({total}, {exp})"
    print(f"({exp} experiments in total)")


# solve using a planner:
def solve_swapple(goal, args):
    pddl_lines = generate_problem_specification(goal)
    # print(pddl_lines)
    # writing pddl to file:
    f = open(args.pddl_problem_out, "w")
    for line in pddl_lines:
        f.write(line)
    f.close()
    run_fdownward(args)
    plan = parse_fdplan(args.plan_file)
    # returning empty plan:
    return plan


# Solve swapple goal with a planner, and print the result
def swapple_main(goal, args):
    start_time = time.perf_counter()
    print(f"\nUsing {args.encoding} encoding,", end=" ")
    print(f"\nUsing Planner {args.planner}")
    plan = solve_swapple(goal, args)
    plan_length = len(plan)
    print_steps(plan, plan_length, args.verbose)
    print(f"\nSolved {sz.N} x {sz.N} matrix in {plan_length} steps")
    total_time = time.perf_counter() - start_time
    print(f"\nTime taken: {total_time}")
    return plan_length


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

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


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


# Using same arguments as swapple_z3.py, easier for integration later and reuse of code:
if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser(
        description="Solve a SWAPPLE instance with a classical planner"
    )
    parser.add_argument(
        "-v", "--verbose", help="show intermediate matrices", action="store_true"
    )
    parser.add_argument(
        "-e",
        "--encoding",
        help="Lifted encoding used by default",
        choices=["lifted"],
        default="lifted",
    )
    parser.add_argument(
        "--aux_files",
        help="location for intermediate files (default ./intermediate_files)",
        default="./intermediate_files",
        metavar="DIR",
    )

    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(
        "--planner",
        help="Any planner from FD (fd-ms)",
        metavar="PLANNER",
        default="fd-ms",
    )
    solver.add_argument(
        "-t",
        "--time",
        type=float,
        help="Solving time limit in seconds, default 1800 seconds",
        default=1800,
    )

    args = parser.parse_args()

    # find Benchmarks and Domains,
    source_location = os.path.dirname(sz.__file__)
    args.domains = os.path.join(source_location, "Domains")

    # we use intermediate directory for intermediate files:
    os.makedirs(args.aux_files, exist_ok=True)
    args.pddl_domain_out = os.path.join(args.domains, "Lifted.pddl")
    args.pddl_problem_out = os.path.join(args.aux_files, "problem.pddl")
    args.log_out = os.path.join(args.aux_files, "log_out")
    args.SAS_file = os.path.join(args.aux_files, "out.sas")
    args.plan_file = os.path.join(args.aux_files, "plan")

    if args.verbose:
        print(args)

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