import os
import sys
import random
import argparse

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

    # Generate probabilistic facts
    for i in range(0, N):
        if use_choices:
            lines.append(f"{{a({i})}}.")
        else:
            lines.append(f"0.5::a({i}).")

    even = [ f"a({i})" for i in range(0, N, 2) ]
    lines.append(f"qr :- {', '.join(even)}.")
    odd = [ f"a({i})" for i in range(1, N, 2) ]
    lines.append(f"qr :- {', '.join(odd)}, not nqr.")
    lines.append(f"nqr :- {', '.join(odd)}, not qr.")
    return lines

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