Evolutionary Seating Arrangement

esa

#!/usr/bin/env python3

import toml
import argparse
from random import randint, sample, choice, shuffle
from collections import namedtuple
from itertools import islice

Guest = namedtuple("Guest", "name, requires, wants")
Table = namedtuple("Table", "capacity, guests")
Arrangement = namedtuple("Arrangement", "tables, guests")
Individual = namedtuple("Individual", "arrangement")
Population = namedtuple("Population", "individuals")
Config = namedtuple("Config", "tables, capacity, population, guests")

def table_swap_guests(ta, tb):
    ga, gb = choice(ta.guests), choice(tb.guests)
    tna = [ g for g in ta.guests if g != ga ]
    tna.append(gb)
    tnb = [ g for g in tb.guests if g != gb ]
    tnb.append(ga)
    return Table(ta.capacity, tna), Table(tb.capacity, tnb)

# XXX: DRY with arr_frag_want
def arr_frag_req(arr, factor=4):
    frag = 0
    for guest in arr.guests.values():
        for table in arr.tables:
            occupants = set([g.name for g in table.guests])
            isec = occupants.intersection(guest.requires)
            frag += 1 if isec else 0
    return frag * factor

# XXX: DRY with arr_frag_req
def arr_frag_want(arr, factor=1):
    frag = 0
    for guest in arr.guests.values():
        for table in arr.tables:
            occupants = set([g.name for g in table.guests])
            isec = occupants.intersection(guest.wants)
            frag += 1 if isec else 0
    return frag * factor

def arr_error(arr):
    return arr_frag_req(arr) + arr_frag_want(arr)

def arr_mutate(arr):
    chosen = sample(arr.tables, 2)
    tables = [ t for t in arr.tables if t not in chosen ]
    tables.extend(table_swap_guests(*chosen))
    return Arrangement(tables, arr.guests)

def ind_create(config):
    assert config.tables * config.capacity >= len(config.guests)
    mingled = list(config.guests.values())
    shuffle(mingled)
    allocated = zip(*[iter(mingled)] * config.capacity)
    tables = [ Table(config.capacity, occupants) for occupants in allocated ]
    return Individual(Arrangement(tables, config.guests))

def ind_splice(a, b):
    # XXX: This isn't splicing, it's just picking the least error
    return a if arr_error(a.arrangement) < arr_error(b.arrangement) else b

def ind_mutate(ind):
    return Individual(arr_mutate(ind.arrangement))

def ind_key(ind):
    return arr_error(ind.arrangement)

def pop_create(config, seed=None):
    pop = list()
    n = config.population

    if seed:
        assert isinstance(seed, list)
        assert config.population >= len(seed)
        n -= len(seed)
        pop.extend(seed)

    for i in range(n):
        pop.append(ind_create(config))

    return pop

def pop_select(pop, n, key=None):
    if not key:
        key = ind_key
    return sorted(pop, key=key)[:n]

def pop_splice(subpop):
    i = iter(subpop)
    return [ ind_splice(a, b) for a, b in zip(*[i] * 2) ]

def pop_mutate(pop):
    selected = randint(1, len(pop) // 2) # at least one
    return [ ind_mutate(ind) for ind in sample(pop, selected) ]

def pop_evolve(config, pop):
    assert 2 < len(pop)
    seed = pop_select(pop, 4) + \
            pop_splice(sample(pop, 4)) + \
            pop_mutate(sample(pop, 4))
    return pop_create(config, seed)

def esa_config(args):
    tc = toml.load(args.config)
    guests = dict()

    # populate guests lookup with guest objects
    for g in tc['guests']:
        assert g['name'] not in guests
        guests[g['name']] = Guest(g['name'], frozenset(g['requires']), frozenset(g['wants']))

    return Config(tc['tables'],
            tc['table']['capacity'],
            args.population,
            guests)

def esa(args):
    config = esa_config(args)
    population = pop_create(config)
    for i in range(args.generations):
        population = pop_evolve(config, population)

    for ind in pop_select(population, 1):
        print("Score: {}".format(ind_key(ind)))
        for i, table in enumerate(ind.arrangement.tables):
            print("Table {}".format(i))
            for guest in table.guests:
                print("\t{}".format(guest.name))
            print()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="party.toml")
    parser.add_argument("--population", type=int, default=10)
    parser.add_argument("generations", type=int, default=30, help="Number of generations")
    return esa(parser.parse_args())

if __name__ == "__main__":
    main()