Canoe polo draw creator

README.md

rod-polo-draw

Generate a good canoe polo draw.

Requirements

Runs on Python 3, with click installed (pip install click).

Usage

$ python solve.py input.json output.json

solve.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  solve.py
#  rod-polo-draw
#

import random
import json
from collections import Counter, namedtuple, defaultdict

import click

Time = namedtuple('Time', 'date game_no')
Team = namedtuple('Team', 'name division')
EmptySlot = namedtuple('EmptySlot', 'time type')
FilledSlot = namedtuple('FilledSlot', 'time type team')
DrawConfig = namedtuple('DrawConfig', 'teams slots divisions')

EmptySlot.fill = lambda self, t: FilledSlot(self.time, self.type, t)


@click.command()
@click.argument('draw_file', type=click.Path(exists=True))
@click.argument('output_file', type=click.Path())
def main(draw_file, output_file):
    "Generate a draw given a JSON draw config."
    config = _load_config(draw_file)
    filled_slots = _solve(_heuristic, config)
    _print_stats(filled_slots)
    _save_draw(filled_slots, output_file)


def _load_config(draw_file):
    data = json.load(open(draw_file))

    teams = _fetch_teams(data)
    divisions = set(t.division for t in teams)
    slots = _generate_empty_slots(data)

    # change the team order arounde each time
    random.shuffle(teams)

    return DrawConfig(teams, slots, divisions)


def _fetch_teams(data):
    teams = []
    for division in data['divisions']:
        for team in division['teams']:
            teams.append(Team(team['name'], team['div']))

    return teams


def _generate_empty_slots(data):
    slots = []
    for block in data['datesAndTimes']['jsonDatesAndTimes']:
        n_games = block['games']
        date = block['date']

        for i in range(1, n_games + 1):
            t = Time(date, i)
            slots.append(EmptySlot(t, 'game'))
            slots.append(EmptySlot(t, 'game'))
            slots.append(EmptySlot(t, 'referee'))

        start = Time(date, 1)
        end = Time(date, n_games)
        slots.append(EmptySlot(Time(date, start), 'duty'))
        slots.append(EmptySlot(Time(date, end), 'duty'))

    return slots


def _solve(heuristic, config):
    filled = []
    for s in config.slots:
        best_team = min(config.teams,
                        key=lambda t: heuristic(t, s, filled, config))
        f = s.fill(best_team)
        filled.append(f)

    return filled


def _heuristic(team, slot, filled, config):
    return sum(
        f(team, slot, filled, config)
        for f in [
            _team_can_only_do_one_thing_at_once,
            _divisions_get_equal_play,
            _teams_get_equal_play,
            _teams_get_equal_refereeing,
            _teams_get_equal_duties,
            _teams_referee_their_own_division,
            _games_have_same_division,
        ]
    )


def _team_can_only_do_one_thing_at_once(team, slot, filled, config):
    same_time = [s.team for s in filled if slot.time == s.time]
    if team in same_time:
        return 1000000

    return 0


def _games_have_same_division(team, slot, filled, config):
    if slot.type == 'game':
        same_time = [s.team for s in filled
                     if slot.time == s.time and s.type == 'game']
        if same_time:
            return 1000000 * (same_time[0].division != team.division)

    return 0


def _divisions_get_equal_play(team, slot, filled, config):
    if slot.type != 'game':
        return 0

    counts = Counter(s.team.division for s in filled if s.type == 'game')
    counts[team.division] += 1
    min_ = min(counts[d] for d in config.divisions)
    max_ = max(counts[d] for d in config.divisions)
    return (max_ - min_) * 200


def _teams_get_equal_play(team, slot, filled, config):
    if slot.type != 'game':
        return 0

    counts = Counter([s.team for s in filled if s.type == 'game'])
    counts[team] += 1
    min_ = min(counts[t] for t in config.teams)
    max_ = max(counts[t] for t in config.teams)
    assert not filled or max_ > 0
    return (max_ - min_) * 1000


def _teams_get_equal_refereeing(team, slot, filled, config):
    if slot.type != 'referee':
        return 0

    counts = Counter([s.team for s in filled if s.type == 'referee'])
    counts[team] += 1
    min_ = min(counts[t] for t in config.teams)
    max_ = max(counts[t] for t in config.teams)
    assert not filled or max_ > 0
    return (max_ - min_) * 10


def _teams_get_equal_duties(team, slot, filled, config):
    if slot.type != 'duty':
        return 0

    counts = Counter([s.team for s in filled if s.type == 'duty'])
    counts[team] += 1
    min_ = min(counts[t] for t in config.teams)
    max_ = max(counts[t] for t in config.teams)
    assert not filled or max_ > 0
    return (max_ - min_) * 10


def _teams_referee_their_own_division(team, slot, filled, config):
    if slot.type == 'referee':
        teams = [s.team for s in filled if s.time == slot.time]
        return 5 * (teams[0].division != team.division)

    return 0


def _is_same_game(s1, s2):
    return s1['date'] == s2['date'] and s1['game_no'] == s2['game_no']


def _print_stats(slots):
    teams = defaultdict(Counter)
    for s in slots:
        teams[s.team][s.type] += 1

    games = [t['game'] for t in teams.values()]
    referees = [t['referee'] for t in teams.values()]
    duty = [t['duty'] for t in teams.values()]

    print('Draw statistics')
    print()
    print('Games:', games)
    print('Referees:', referees)
    print('Duties:', duty)


def _save_draw(slots, output_file):
    draw = []
    for s in slots:
        draw.append({
            'date': s.time.date,
            'game_no': s.time.game_no,
            'type': s.type,
            'team': {
                'name': s.team.name,
                'div': s.team.division,
            },
        })

    json.dump(draw, open(output_file, 'w'), indent=2)


if __name__ == '__main__':
    main()