import os
import sys
import random
import argparse

def create_hmm_program(N, add_comments, use_eclauses, use_choices, nondisjunctive):
    lines = []
    if add_comments:
        lines.append(f"% Size: {N}")
    # Print probabilistic atoms in reverse order
    if use_choices:
        # Start with manifest variables probabilities
        for i in range(N, 0, -1):
            lines.append("{b(" + str(i) + ")}.")
        # Then write latent variable probabilities
        for i in range(N, 0, -1):
            lines.append("{a(" + str(i) + ")}.")
    else:
        # Start with manifest variables probabilities
        for i in range(N, 0, -1):
            lines.append(f"0.5::b({i}).")
        # Then write latent variable probabilities
        for i in range(N, 0, -1):
            lines.append(f"0.5::a({i}).")
    # Now specify rule of chain
    for i in range(N, 1, -1):
        lines.append(f"x({i}) :- a({i}), x({i-1}).")
    lines.append("x(1) :- a(1).")
    # Finally specify manifest variables
    if nondisjunctive:
        for i in range(N, 0, -1):
            lines.append(f"y({i},0) :- x({i}), not y({i},1).")
            lines.append(f"y({i},1) :- x({i}), not y({i},0).")
            lines.append(f"y({i},0) :- b({i}), not x({i}).")
            lines.append(f"y({i},1) :- not b({i}), not x({i}).")
    else:
        for i in range(N, 0, -1):
            lines.append(f"y({i},0); y({i},1) :- x({i}).")
            lines.append(f"y({i},0) :- b({i}), not x({i}).")
            lines.append(f"y({i},1) :- not b({i}), not x({i}).")
    return lines

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Create HMM-like programs.')
    parser.add_argument('beginning', type=int, help='Start size of chain')
    parser.add_argument('end', type=int, nargs='?', default=None, help='End size of chain')
    parser.add_argument('-c', action='store_true', help='Add comments to output')
    parser.add_argument('-e', action='store_true', help='Encode using exactly-one clauses')
    parser.add_argument('-nd', action='store_true', help='Encode using normal rules (non-disjunctive)')
    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_eclauses = args.e
    use_choices = args.choice
    nondisjunctive = args.nd

    if use_eclauses:
        print('%%% Using exactly-one clauses')
    if nondisjunctive:
        print('%%% Using non-disjunctive rules')
    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/hmm{N}"
        os.makedirs(directory, exist_ok=True)
        file_path = os.path.join(directory, f"hmm{N}.pasp")
        with open(file_path, 'w') as f:
            lines = create_hmm_program(N, add_comments, use_eclauses, use_choices, nondisjunctive)
            f.write("\n".join(lines))
