import os
import sys
import argparse
import math

def create_ground_food_program_seq(n, m, add_comments=True):
    """
    Generates the ASP program for the food preference problem using sequential counters
    and annotated disjunctions.

    Args:
        n: The number of people.
        m: The number of food items.
        probabilities: A list of m probabilities for food preferences.
                       If None, uniform probabilities (1/m) are assumed.
        add_comments: Whether to add comments to the output.

    Returns:
        A list of strings, where each string is a line of the ASP program.
    """
    lines = []
    if add_comments:
        lines.append(f"% Grounded Food Preference Problem with n={n}, m={m}")
        lines.append(f"% Minimum people choosing a food for it to be valid: floor({n}/2) + 1 = {math.floor(n/2) + 1}")
        lines.append("% Constraints are unrolled.")
        lines.append("")

    # Generate grounded basic facts
    if add_comments:
        lines.append("% Basic facts")
    for i in range(1, n + 1):
        lines.append(f"1.0::person({i}).")
    for j in range(1, m + 1):
        lines.append(f"1.0::food({j}).")
    lines.append("")

    # Unroll the constraint: For each person, exactly one food preference is chosen.
    # Assuming an underlying constraint like: 1 { person_food_preference(P, F) : food(F) } 1 :- person(P).
    if add_comments:
        lines.append("% Unrolled constraint: For each person, exactly one food preference is true.")
    for i in range(1, n + 1):
        # At least one preference for person i
        #at_least_one_body = ", ".join([f"not person_food_preference({i},{j})" for j in range(1, m + 1)])
        #if m > 0: # Only add if there's at least one food
        #    lines.append(f":- {at_least_one_body}.")

        # At most one preference for person i (in this case, it results in exactly one)
        for j in range(1, m + 1):
            at_most_one_body = ", ".join([f"not person_food_preference({i},{k})" for k in range(1, m + 1) if k != j])
            lines.append(f"person_food_preference({i},{j}) :- {at_most_one_body}.")
    lines.append("")

    # --- Chooses Rule ---
    # Generate grounded chooses rules
    if add_comments:
        lines.append("% Grounded chooses rules")
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            # Since person(i) and food(j) are facts, the rule simplifies
            lines.append(f"chooses({i},{j}) :- person_food_preference({i},{j}).")
    lines.append("")

    # Calculate the minimum number of people required for a valid food
    min_people = math.floor(n / 2) + 1

    # --- Valid Food (Sequential Counter for "at least min_choosers") ---
    if add_comments:
        lines.append(f"% Valid Food: At least {min_people} people must choose it (encoded with sequential counter)")
        lines.append(f"% vfc_count(FoodID, PersonIDMax, NumChoosing): NumChoosing people from 1..PersonIDMax choose FoodID")

    for f_id in range(1, m + 1):
        # Forces vfc_count to be true for food f_id
        #lines.append(f"vfc_count({f_id}, 0, 0) :- vfc_count({f_id}, 0, 0).") # Base case: for food f_id, after 0 people, 0 chosen.
        for p_id in range(1, n + 1):
            for k_prev in range(p_id): # Max k_prev can be p_id-1 (if all previous P-1 people chose it)
                k_curr_if_chosen = k_prev + 1
                # If person p_id chooses food f_id
                if k_prev > 0:
                    lines.append(f"vfc_count({f_id}, {p_id}, {k_curr_if_chosen}) :- vfc_count({f_id}, {p_id-1}, {k_prev}), chooses({p_id},{f_id}).")
                    # If person p_id does NOT choose food f_id
                    lines.append(f"vfc_count({f_id}, {p_id}, {k_prev}) :- vfc_count({f_id}, {p_id-1}, {k_prev}), not chooses({p_id},{f_id}).")
                else:
                    # If person p_id chooses food f_id
                    lines.append(f"vfc_count({f_id}, {p_id}, {k_curr_if_chosen}) :- chooses({p_id},{f_id}).")
                    # If person p_id does NOT choose food f_id
                    lines.append(f"vfc_count({f_id}, {p_id}, 0) :- not chooses({p_id},{f_id}).")
    lines.append("")
    if add_comments:
        lines.append(f"% valid_food(F) definition using the counter")
    for f_id in range(1, m + 1):
        for k_final in range(int(min_people), n + 1): # If K_final is min_people or more
            lines.append(f"valid_food({f_id}) :- vfc_count({f_id}, {n}, {k_final}).")
    lines.append("")

    # --- Chosen Food: Exactly one food from valid_food(F) (Simplified Encoding) ---
   # Define chosen_food based on valid_food and absence of other chosen valid foods
    for j in range(1, m + 1):
        # At least one valid food
        # at_least_one_body = ", ".join([f"not valid_food({j})" for j in range(1, m + 1)])
        # if m > 0: # Only add if there's at least one food
        #     lines.append(f":- {at_least_one_body}.")

        # At most one chosen food (in this case, it results in exactly one)
        # Generate the body of the rule with chosen_food for each valid food
        at_most_one_body = ", ".join([f"not chosen_food({k})" for k in range(1, m + 1) if k != j])
        lines.append(f"chosen_food({j}) :- valid_food({j}), {at_most_one_body}.")

    return lines

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Create Grounded Food Preference Problem programs.')
    parser.add_argument('n_people_start', type=int, help='Number of people (n)')
    parser.add_argument('n_people_end', type=int, help='Number of people (n)')
    parser.add_argument('m_food', type=int, help='Number of food items (m)')
    parser.add_argument('-c', action='store_true', help='Add comments to output')

    args = parser.parse_args()

    beginning = args.n_people_start
    end = args.n_people_end if args.n_people_end is not None else beginning
    m_foods = args.m_food
    add_comments = args.c

    if beginning < 3:
        print("Error: The number of people must be at least 3.")
        parser.print_help()
        exit(1)

    if m_foods < 3:
        print("Error: The number of food items must be at least 4.")
        parser.print_help()
        exit(1)

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