fulvia/modules/dice.py

#!/usr/bin/env python3
"""
Dice rolling, the core function of any IRC bot.
"""
import random
import re
import operator

import module
from tools.calculation import eval_equation


class DicePouch:
	def __init__(self, num_of_die, type_of_die, addition):
		"""
		Initialize dice pouch and roll the dice.

		Args:
			num_of_die: number of dice in the pouch.
			type_of_die: how many faces the dice have.
			addition: how much is added to the result of the dice.
		"""
		self.num = num_of_die
		self.type = type_of_die
		self.addition = addition

		self.dice = {}
		self.dropped = {}

		self.roll_dice()

	def roll_dice(self):
		"""Roll all the dice in the pouch."""
		self.dice = {}
		self.dropped = {}
		for __ in range(self.num):
			number = random.randint(1, self.type)
			count = self.dice.setdefault(number, 0)
			self.dice[number] = count + 1

	def drop_lowest(self, n):
		"""
		Drop n lowest dice from the result.

		Args:
			n: the number of dice to drop.
		"""

		sorted_x = sorted(self.dice.items(), key=operator.itemgetter(0))

		for i, count in sorted_x:
			count = self.dice[i]
			if n == 0:
				break
			elif n < count:
				self.dice[i] = count - n
				self.dropped[i] = n
				break
			else:
				self.dice[i] = 0
				self.dropped[i] = count
				n = n - count

		for i, count in self.dropped.items():
			if self.dice[i] == 0:
				del self.dice[i]

	def get_simple_string(self):
		"""Return the values of the dice like (2+2+2[+1+1])+1."""
		dice = self.dice.items()
		faces = ("+".join([str(face)] * times) for face, times in dice)
		dice_str = "+".join(faces)

		dropped_str = ""
		if self.dropped:
			dropped = self.dropped.items()
			dfaces = ("+".join([str(face)] * times) for face, times in dropped)
			dropped_str = "[+%s]" % ("+".join(dfaces),)

		plus_str = ""
		if self.addition:
			plus_str = "{:+d}".format(self.addition)

		return "(%s%s)%s" % (dice_str, dropped_str, plus_str)

	def get_compressed_string(self):
		"""Return the values of the dice like (3x2[+2x1])+1."""
		dice = self.dice.items()
		faces = ("%dx%d" % (times, face) for face, times in dice)
		dice_str = "+".join(faces)

		dropped_str = ""
		if self.dropped:
			dropped = self.dropped.items()
			dfaces = ("%dx%d" % (times, face) for face, times in dropped)
			dropped_str = "[+%s]" % ("+".join(dfaces),)

		plus_str = ""
		if self.addition:
			plus_str = "{:+d}".format(self.addition)

		return "(%s%s)%s" % (dice_str, dropped_str, plus_str)

	def get_sum(self):
		"""Get the sum of non-dropped dice and the addition."""
		result = self.addition
		for face, times in self.dice.items():
			result += face * times
		return result

	def get_number_of_faces(self):
		"""
		Returns sum of different faces for dropped and not dropped dice

		This can be used to estimate, whether the result can be shown in
		compressed form in a reasonable amount of space.
		"""
		return len(self.dice) + len(self.dropped)


def _roll_dice(bot, dice_expression):
	result = re.search(
		r"""
		(?P<dice_num>-?\d*)
		d
		(?P<dice_type>-?\d+)
		(v(?P<drop_lowest>-?\d+))?
		$""",
		dice_expression,
		re.IGNORECASE | re.VERBOSE)

	dice_num = int(result.group('dice_num') or 1)
	dice_type = int(result.group('dice_type'))

	# Dice can't have zero or a negative number of sides.
	if dice_type <= 0:
		bot.reply("I don't have any dice with %d sides. =(" % dice_type)
		return None  # Signal there was a problem

	# Can't roll a negative number of dice.
	if dice_num < 0:
		bot.reply("I'd rather not roll a negative amount of dice. =(")
		return None  # Signal there was a problem

	# Upper limit for dice should be at most a million. Creating a dict with
	# more than a million elements already takes a noticeable amount of time
	# on a fast computer and ~55kB of memory.
	if dice_num > 1000:
		bot.reply('I only have 1000 dice. =(')
		return None  # Signal there was a problem

	dice = DicePouch(dice_num, dice_type, 0)

	if result.group('drop_lowest'):
		drop = int(result.group('drop_lowest'))
		if drop >= 0:
			dice.drop_lowest(drop)
		else:
			bot.reply("I can't drop the lowest %d dice. =(" % drop)

	return dice


@module.commands("roll", "dice", "d")
@module.example(".roll 3d1+1", "You roll 3d1+1: (1+1+1)+1 = 4")
def roll(bot, trigger):
	"""
	.dice XdY[vZ][+N], rolls dice and reports the result.

	X is the number of dice. Y is the number of faces in the dice. Z is the
	number of lowest dice to be dropped from the result. N is the constant to
	be applied to the end result.
	"""
	# This regexp is only allowed to have one captured group, because having
	# more would alter the output of re.findall.
	dice_regexp = r"-?\d*[dD]-?\d+(?:[vV]-?\d+)?"

	# Get a list of all dice expressions, evaluate them and then replace the
	# expressions in the original string with the results. Replacing is done
	# using string formatting, so %-characters must be escaped.
	if len(trigger.args) < 2:
		return bot.reply("No dice to roll.")
	arg_str = ' '.join(trigger.args[1:])
	dice_expressions = re.findall(dice_regexp, arg_str)
	arg_str = arg_str.replace("%", "%%")
	arg_str = re.sub(dice_regexp, "%s", arg_str)

	f = lambda dice_expr: _roll_dice(bot, dice_expr)
	dice = list(map(f, dice_expressions))

	if None in dice:
		# Stop computing roll if there was a problem rolling dice.
		return

	def _get_eval_str(dice):
		return "(%d)" % (dice.get_sum(),)

	def _get_pretty_str(dice):
		if dice.num <= 10:
			return dice.get_simple_string()
		elif dice.get_number_of_faces() <= 10:
			return dice.get_compressed_string()
		else:
			return "(...)"

	eval_str = arg_str % (tuple(map(_get_eval_str, dice)))
	pretty_str = arg_str % (tuple(map(_get_pretty_str, dice)))

	# Showing the actual error will hopefully give a better hint of what is
	# wrong with the syntax than a generic error message.
	try:
		result = eval_equation(eval_str)
	except Exception as e:
		bot.reply("SyntaxError, eval(%s), %s" % (eval_str, e))
		return

	bot.reply("You roll %s: %s = %d" % (
		' '.join(trigger.args[1:]), pretty_str, result))


@module.commands("choice", "choose")
@module.example(".choose opt1,opt2,opt3")
def choose(bot, trigger):
	"""
	.choice option1|option2|option3 - Makes a difficult choice easy.
	"""
	if len(trigger.args) < 2:
		return bot.reply("I'd choose an option, but you didn't give me any.")
	msg = ' '.join(trigger.args[1:])
	choices = [msg]
	for delim in '|\\/,':
		choices = msg.split(delim)
		if len(choices) > 1:
			break
	# Use a different delimiter in the output, to prevent ambiguity.
	for show_delim in ',|/\\':
		if show_delim not in msg:
			show_delim += ' '
			break

	pick = random.choice(choices)
	msg = f"Your options: {show_delim.join(choices)}. My choice: {pick}"
	bot.reply(msg)