import os
import sys
import random
import argparse

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

    # Generate probabilistic facts for queens
    for r in range(1, N + 1):
        for c in range(1, N + 1):
            p = random.random()
            if use_choices:
                lines.append(f"{{queen({r},{c})}}.")
            else:
                lines.append(f"{p:.2f}::queen({r},{c}).")

    # Generate conflict rules for rows
    for r in range(1, N + 1):
        for c1 in range(1, N):
            for c2 in range(c1 + 1, N + 1):
                lines.append(f"conflict :- queen({r},{c1}), queen({r},{c2}).")

    # Generate conflict rules for columns
    for c in range(1, N + 1):
        for r1 in range(1, N):
            for r2 in range(r1 + 1, N + 1):
                lines.append(f"conflict :- queen({r1},{c}), queen({r2},{c}).")

    # Generate conflict rules for diagonals
    for r1 in range(1, N):
        for c1 in range(1, N):
            for r2 in range(r1 + 1, N + 1):
                for c2 in range(1, N + 1):
                    if abs(r1 - r2) == abs(c1 - c2):
                        lines.append(f"conflict :- queen({r1},{c1}), queen({r2},{c2}).")

    return lines

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Create N-Queens Problem programs.')
    parser.add_argument('beginning', type=int, help='Start size of board')
    parser.add_argument('end', type=int, nargs='?', default=None, help='End size of board (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/queens{N}"
        os.makedirs(directory, exist_ok=True)
        file_path = os.path.join(directory, f"queens{N}.pasp")
        with open(file_path, 'w') as f:
            lines = create_queens_program(N, add_comments, use_choices)
            f.write("\n".join(lines))
