Converts a CSV of availabilities into Prolog code that can make assignments

csv_to_prolog.py

# csv_to_prolog.py
#
# Takes a CSV from stdin in the following format:
#   <Any>, Time1, Time2, Time3, Time4...
#  counts,     2,     3,     4,     5... (number of spots for each time)
# Person1,   Yes, Maybe,   Yes,    No... (Yes, Maybe, or No for availability)
# Person2,   Yes, Maybe,   Yes,    No...
#
# and outputs Prolog to do assignments.
#
# Example usage for bash (assuming SWI Prolog is installed and on path):
#   swipl -s <(cat input.csv | python3 csv_to_prolog.py) -g run -t halt
# and for fish:
#   swipl -s (cat input.csv | python3 csv_to_prolog.py | psub) -g run -t halt
#
# Replace `run` with `run($x)` to allow $x assignments to Maybe availability.

import sys

prolog_body = """
/* Code */
last_count(_Type, 0, []).
last_count(Type, Count, [[Type, Count]|_Rest]).
last_count(Type, Count, [[OtherType, _OtherCount]|Rest]) :-
  Type \== OtherType,
  last_count(Type, Count, Rest).

assign_helper([], [], 0).

assign_helper([Person|PeopleRest], [[Time, Count]|Rest], 0) :-
  assign_helper(PeopleRest, Rest, 0),
  available(Person, Time, 'Yes'),
  last_count(Time, LastCount, Rest),
  Count is LastCount + 1,
  spots_at_time(Time, MaxCount),
  Count =< MaxCount.

assign_helper([Person|PeopleRest], [[Time, Count]|Rest], MaybeCount) :-
  assign_helper(PeopleRest, Rest, LastMaybeCount),
  available(Person, Time, 'Maybe'),
  last_count(Time, LastCount, Rest),
  Count is LastCount + 1,
  MaybeCount is LastMaybeCount + 1,
  spots_at_time(Time, MaxCount),
  Count =< MaxCount.

times_process([], [], []).
times_process([P|PRest], [[T, _C]|TCRest], [[P, T]|Rest]) :-
  times_process(PRest, TCRest, Rest).

assign(Assignments, MaxMaybes) :-
  bagof(E, person(E), People),
  assign_helper(People, TimesCounts, MaybeCount),
  MaybeCount =< MaxMaybes,
  times_process(People, TimesCounts, Assignments).

assign(Assignments) :- assign(Assignments, 0).

display([Person, Time]) :-
  write(Person),
  write(", "),
  writeln(Time).

run :- run(0).

run(MaxMaybes) :-
  assign(Assignments, MaxMaybes),
  maplist(display, Assignments).
"""

print('/* Generated by csv_to_prolog.py */')
print()

print('/* Data */')
times = None
counts = None
people = []
for line in sys.stdin:
    line = line.strip()
    parts = [x.strip() for x in line.split(',')]
    if not times:
        times = parts[1:]
        continue
    name = parts[0]
    if not name:
        continue
    if name == 'counts':
        counts = [int(c) for c in parts[1:]]
        for time, count in zip(times, counts):
            print("spots_at_time('{}', {}).".format(time, count))
        continue
    people.append((name, parts[1:], 
        sum([5*c if x == 'Yes' else \
            (1*c if x == 'Maybe' else 0) for c, x in zip(counts, parts[1:])])))

if len(people) > sum(counts):
    print("Warning! {} spots for {} people".format(sum(counts), len(people)), \
          file=sys.stderr)

people = sorted(people, key=lambda x: -x[2]) # least constrained first

if len(sys.argv) > 1:
    people = people[:int(sys.argv[1])]

for name, _, _ in people:
    print("person('{}').".format(name))

for name, prefs, score in people:
    how_many = 0
    for time, pref in zip(times, prefs):
        if pref == 'No':
            continue
        assert(pref in ['Yes', 'Maybe'])
        how_many += 1
        print("available('{}', '{}', '{}').".format(name, time, pref))
    if not how_many:
        print("Warning! {} has no availability".format(name), file=sys.stderr)

#print('all_people([{}]).'.format(', '.join(people)))

print(prolog_body)

example-availabilities.csv