Toy program for playing with D'Hondt system

dhontd.py

import argparse
import logging
import random
import string

from itertools import groupby
from operator import itemgetter

PARTIES = ["Kesk", "Kok", "SDP", "Sin", "PS", "Vihr", "Vas", "RKP", "KD", "TL", "EOP", "FP", "IPU", "KP", "KTP", "Lib",
           "PP", "SKE", "SKP"]

logging.basicConfig(level=logging.INFO)


def main():
    def split_at(lst, index):
        return lst[:index + 1], lst[index + 1:]

    def chunks(l, n):
        for i in range(0, len(l), n):
            yield l[i:i + n]

    def print_votes(lst, max_rows=None):
        for i, chunk in enumerate(chunks(lst, 5), start=1):
            if max_rows and i > max_rows:
                break

            for party, idx, votes in chunk:
                print('{}{}: {:6d}'.format(party, idx, int(votes)), end="\t")
            print()
        print()

    def rank_order(elect_count, votes):
        def gen_weighted_votes(label, votes, max_count):
            return [(label, i + 1, votes / float(i + 1)) for i in range(max_count)]

        weighted_votes = []
        for k, v in votes.items():
            weighted_votes.extend(gen_weighted_votes(k, v, 10 * elect_count))

        return weighted_votes

    parser = argparse.ArgumentParser(description='Show D\'Hondt vote distribution')
    parser.add_argument('votes', metavar='N', type=int, nargs='*',
                        help='Votes a party gets. First for A, second for B, third for C etc')
    parser.add_argument('--seats', '-s', type=int, help='Number of seats to elect in total')
    parser.add_argument('--party-count', '-p', type=int, default=4)
    parser.add_argument('--num-voters', '-v', type=int, default=int(3E6), help='Total number of voters (use with -f)')
    parser.add_argument('--max-party-votes', '-m', type=float, default=30.0,
                        help='Max number of votes for party (percentage)')
    parser.add_argument('--fun-mode', '-f', action='store_true',
                        help='Map abstract party symbols to real parties, implies -c 19')
    args = parser.parse_args()

    zipper = lambda i: dict(zip(range(len(i)), i))

    pos_to_party = zipper(string.ascii_uppercase)
    if args.fun_mode:
        logging.info("Funny mode enabled")
        args.party_count = len(PARTIES)
        args.seats = args.seats or 200
        random.shuffle(PARTIES)
        pos_to_party = zipper(PARTIES)
    else:
        args.seats = args.seats or 3

    vote_list = args.votes
    if not vote_list:
        print(f'Using randomized votes for {args.party_count} parties')
        vote_list = []

        voters_left = args.num_voters
        for _ in range(args.party_count):
            vote_max_limit = min((args.max_party_votes / 100.0) * args.num_voters, voters_left)
            votes = random.randint(0, vote_max_limit)
            vote_list.append(votes)
            voters_left -= votes

        vote_list = sorted(vote_list, reverse=True)

        for i, chunk in enumerate(chunks(vote_list, 5)):
            print(', '.join(['{}: {}'.format(pos_to_party[5 * i + j], v) for j, v in enumerate(chunk)]))

    vote_map = {}
    for i, num_votes in enumerate(vote_list):
        label = pos_to_party[i]
        vote_map[label] = num_votes

    total_votes = sum(vote_list)
    elected, rest = split_at(sorted(rank_order(args.seats, vote_map), key=lambda x: -x[-1]), args.seats - 1)

    print(f'Total votes: {total_votes}')
    print()

    print(f'Elected {args.seats} people, ordered by relative votes:')
    print_votes(elected)

    sorted_by_party = sorted(elected, key=itemgetter(0))

    print('Seats and proportionate seats:')
    by_seats = []
    for party, entries in groupby(sorted_by_party, key=itemgetter(0)):
        values = list(entries)
        seats = len(values)
        top_seats = values[0][-1]
        prop = args.seats * top_seats / total_votes
        by_seats.append((seats, party, prop))

    for seats, party, prop in sorted(by_seats, reverse=True):
        print('{:>6s}: {:<6d} / {:.2f}'.format(party, seats, prop))

    print()
    print("Top 3 rows for not chosen, ordered by relative votes:")
    print_votes(rest, max_rows=3)


main()