import os
import argparse

def append_rules(lines, depth, N, inside, default_prefix="infected"):
    aux = "_" if inside else ""
    prefix = default_prefix + aux + "_".join(map(str, inside))

    if depth == N:
        for i in range(1, N + 1):
            if i not in inside:
                lines.append(f"{prefix}({i}) :- contaminated({i}).")
        return

    for i in range(1, N + 1):
        if i not in inside:
            lines.append(f"{prefix}({i}) :- contaminated({i}).")
            for j in range(1, N + 1):
                if (j != i) and (j not in inside):
                    new_list = sorted(inside + [i])
                    sub_prefix = f"{default_prefix}_" + "_".join(map(str, new_list))
                    lines.append(f"{prefix}({i}) :- friend({i},{j}), {sub_prefix}({j}).")
            if (len(inside) == 0) or (inside[-1] < i):
                append_rules(lines, depth + 1, N, sorted(inside + [i]), default_prefix)

def create_pin_program(N, add_comments, use_choices):
    lines = []
    if add_comments:
        lines.append(f"% Size: {N}")

    # Add probabilistic facts for contaminated atoms
    for i in range(1, N + 1):
        if use_choices:
            lines.append(f"{{contaminated({i})}}.")
        else:
            lines.append(f"0.5::contaminated({i}).")

    # Add probabilistic facts for friend atoms
    for i in range(1, N + 1):
        for j in range(1, N + 1):
            if i != j:
                if use_choices:
                    lines.append(f"{{friend({i},{j})}}.")
                else:
                    lines.append(f"0.5::friend({i},{j}).")

    # Add logical rules for observed atoms
    # They should be of form:
        # infected(i) :- contaminated(i), not healthy(i).
        # infected(i) :- friend(i,j), contaminated(j), not healthy(i).
    # But we do positive cycle-breaking
    append_rules(lines, 1, N, [], default_prefix="infected")

    for i in range(1, N + 1):
        lines.append(f"healthy({i}) :- not infected({i}).")
        # Add non-stratified negation
        lines.append(f"symptomatic({i}) :- infected({i}), not vector({i}).")
        lines.append(f"vector({i}) :- infected({i}), not symptomatic({i}).")

    return lines

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Create Probabilistic Interaction Network programs.')
    parser.add_argument('beginning', type=int, help='Start size of network')
    parser.add_argument('end', type=int, nargs='?', default=None, help='End size of network (optional, defaults to beginning)')
    parser.add_argument('-c', action='store_true', help='Add comments to output')
    parser.add_argument('-choice', action='store_true', help='Encode probabilistic facts as choice rules')

    args = parser.parse_args()

    beginning = args.beginning
    end = args.end if args.end is not None else beginning
    add_comments = args.c
    use_choices = args.choice

    if use_choices:
        print('%%% Using choice rules instead of probabilistic facts')

    if beginning == 0 or end == 0 or beginning > end:
        parser.print_help()
        exit(1)

    for N in range(beginning, end + 1):
        directory = f"plp/programs/pin{N}"
        os.makedirs(directory, exist_ok=True)
        file_path = os.path.join(directory, f"pin{N}.pasp")
        with open(file_path, 'w') as f:
            lines = create_pin_program(N, add_comments, use_choices)
            f.write("\n".join(lines))
