updated randomized roles script

#!/usr/bin/env python3

import random
import sys

REQUIRED_ROLES = ['wolf']
WOLF_ROLES = ['wolf', 'werecrow', 'wolf cub', 'traitor', 'hag', 'sorcerer']
NEUTRAL_ROLES = ['fool', 'monster', 'crazed shaman', 'vengeful ghost', 'clone', 'minion', 'cultist'] #minion and cultist aren't really neutral, but work better here
SAFE_ROLES = ['village drunk', 'harlot', 'matchmaker', 'hunter', 'shaman']
VILLAGER_ROLES = ['seer', 'oracle', 'augur', 'detective', 'village drunk',
                  'harlot', 'guardian angel', 'bodyguard', 'village elder',
                  'time lord', 'mad scientist', 'matchmaker', 'hunter', 'shaman',
                  'lycan']

TEMPLATES = ['cursed villager', 'gunner', 'mayor', 'assassin', 'amnesiac', 'bureaucrat']

players = int(sys.argv[1])


roles = REQUIRED_ROLES
wolves = int(random.randint(25,35)/100*players+.5)-1
safes = int(random.randint(15,45)/100*players+.5)-1
neutrals = int(random.randint(0,15)/100*players+.5)
templates = int(random.randint(5,40)/100*players+.5)

roles += [random.choice(WOLF_ROLES) for i in range(wolves)]
roles += [random.choice(NEUTRAL_ROLES) for i in range(neutrals)]
roles += [random.choice(SAFE_ROLES)] + [random.choice(VILLAGER_ROLES) for i in range(safes)]

cursed = False
gunner = False
amnesiac = False

for role in roles:
    if role not in ('seer', 'oracle', 'wolf', 'werecrow', 'wolf cub', 'mad scientist', 'monster', 'fool'):
        cursed = True

    if role not in ('wolf', 'werecrow', 'wolf cub', 'cultist', 'minion',
                    'traitor', 'hag', 'sorcerer', 'fool', 'lycan'):
        gunner = True

    if role not in ('wolf', 'werecrow', 'wolf cub', 'villager',
                    'village drunk', 'village elder', 'time lord', 'cultist',
                    'minion', 'traitor', 'lycan', 'vengeful ghost', 'clone',
                    'fool', 'monster', 'matchmaker'):
        amnesiac = True

if cursed and "seer" in roles or "oracle" in roles:
    TEMPLATES.remove('cursed villager')

if not gunner:
    TEMPLATES.remove('gunner')

if not amnesiac:
    TEMPLATES.remove('amnesiac')

roles += [random.choice(TEMPLATES) for i in range(templates)]

if 'gunner' in roles:
    roles.append('sharpshooter')

print('!fgame roles=' + ','.join(role + ':1' for role in roles))