bot/services/roll20.py

import random

import parsy


# TODO: functions
# TODO: nested groups


ROLL_LIMIT = 1000


def num(x):
    if int(x) == x:
        return int(x)
    return float(x)


class Group(object):
    def __init__(self, items):
        self.items = items
        self.keep = None
        self.drop = None
        self.succ = None
        self.fail = None

    def __repr__(self):
        return '<group {0} keep:{1} drop:{2} succ:{3} fail:{4}>'.format(
            ' '.join([repr(x) for x in self.items]), self.keep, self.drop, self.succ, self.fail)

    def __str__(self):
        calculated = self._calculated()
        kept = self.kept()
        result = []
        for i, x in enumerate(self.items):
            if i in kept:
                if len(self.items) > 1 and self.succ:
                    if self.succ(calculated[i]):
                        result.append('__{0}__'.format(x))
                    elif self.fail and self.fail(calculated[i]):
                        result.append('*{0}*'.format(x))
                    else:
                        result.append(str(x))
                else:
                    result.append(str(x))
            else:
                result.append('~~*{0}*~~'.format(x))
        return '**{{** {0} **}}**'.format(' + '.join(result))

    def _subrolls(self, tree):
        def traverse(node, subrolls):
            try:
                traverse(node.left, subrolls)
                traverse(node.right, subrolls)
            except AttributeError:
                try:
                    node.result
                    subrolls.append(node)
                except AttributeError:
                    return
        subrolls = []
        traverse(tree, subrolls)
        return subrolls

    def _make_expression(self, tree):
        def exp(x):
            def copy(node):
                try:
                    return Operation(node.op, node.func, copy(node.left), copy(node.right))
                except AttributeError:
                    try:
                        node.result
                        return num(x)
                    except AttributeError:
                        return num(node)
            new_tree = copy(tree)
            try:
                return new_tree.calc()
            except AttributeError:
                return num(new_tree)
        return exp

    def _update_subrolls(self, subrolls):
        results = [(r, i, j, s) for i, s in enumerate(subrolls) for j, r in enumerate(s.result) if j in s.kept()]
        results = sorted(results, key=lambda x: (x[0], len(subrolls) - x[1], len(x[3].result) - x[2]))
        if self.keep:
            if self.keep[1]:
                results = results[:round(self.keep[0])]
            else:
                results = results[-round(self.keep[0]):]
        if self.drop:
            if self.drop[1]:
                results = results[:-round(self.drop[0])]
            else:
                results = results[round(self.drop[0]):]
        for subroll in subrolls:
            subroll.group_kept = []
        for _, _, i, subroll in results:
            subroll.group_kept.append(i)

    def _calculated(self):
        # FIXME: nested groups
        if len(self.items) == 1:
            subrolls = self._subrolls(self.items[0])
            self._update_subrolls(subrolls)
            if len(subrolls) == 1:
                exp = self._make_expression(self.items[0])
                if self.succ:
                    subrolls[0].succ = lambda x: self.succ(exp(x))
                    if self.fail:
                        subrolls[0].fail = lambda x: self.fail(exp(x))
                return [subrolls[0].calc()]
            elif self.succ:
                raise RuntimeError('Multiple rolls in a single subroll are not allowed!')
        result = []
        for item in self.items:
            try:
                result.append(item.calc())
            except AttributeError:
                result.append(num(item))
        return result

    def kept(self):
        if len(self.items) == 1:
            return [0]
        calculated = self._calculated()
        result = sorted(enumerate(calculated), key=lambda x: (x[1], len(calculated) - x[0]))
        if self.keep:
            if self.keep[1]:
                result = result[:round(self.keep[0])]
            else:
                result = result[-round(self.keep[0]):]
        if self.drop:
            if self.drop[1]:
                result = result[:-round(self.drop[0])]
            else:
                result = result[round(self.drop[0]):]
        return list(list(zip(*sorted(result)))[0])

    def filtered(self):
        calculated = self._calculated()
        kept = self.kept()
        return [x for i, x in enumerate(calculated) if i in kept]

    def calc(self):
        filtered = self.filtered()
        if len(self.items) > 1:
            if self.succ:
                result = len([x for x in filtered if self.succ(x)])
                if self.fail:
                    result -= len([x for x in filtered if self.fail(x)])
                return result
        return sum(filtered)


class Operation(object):
    def __init__(self, op, func, left, right):
        self.op = op
        self.func = func
        self.left = left
        self.right = right

    def __repr__(self):
        return '<{0} {1} {2}>'.format(self.op, repr(self.left), repr(self.right))

    def __str__(self):
        # FIXME: get rid of unneeded parentheses
        return '( {0} {1} {2} )'.format(str(self.left), self.op.replace('*', '\\*'), str(self.right))

    def calc(self):
        try:
            left = self.left.calc()
        except AttributeError:
            left = num(self.left)
        try:
            right = self.right.calc()
        except AttributeError:
            right = num(self.right)
        return self.func(left, right)


class Roll(object):
    def __init__(self, result):
        self.result = result
        self.label = None
        self.keep = None
        self.drop = None
        self.succ = None
        self.fail = None
        self.group_kept = None

    def __repr__(self):
        return '<roll {0} label:{1} keep:{2} drop:{3} succ:{4} fail:{5}>'.format(
            self.result, self.label, self.keep, self.drop, self.succ, self.fail)

    def __str__(self):
        kept = self.kept()
        result = []
        for i, x in enumerate(self.result):
            if i in kept:
                if self.succ:
                    if self.succ(x):
                        result.append('__{0}__'.format(x))
                    elif self.fail and self.fail(x):
                        result.append('*{0}*'.format(x))
                    else:
                        result.append(str(x))
                else:
                    result.append(str(x))
            else:
                result.append('~~*{0}*~~'.format(x))
        return '**(** {0} **)**'.format(' + '.join(result))

    def kept(self):
        if self.group_kept is not None:
            return self.group_kept
        result = sorted(enumerate(self.result), key=lambda x: (x[1], len(self.result) - x[0]))
        if self.keep:
            if self.keep[1]:
                result = result[:round(self.keep[0])]
            else:
                result = result[-round(self.keep[0]):]
        if self.drop:
            if self.drop[1]:
                result = result[:-round(self.drop[0])]
            else:
                result = result[round(self.drop[0]):]
        return list(list(zip(*sorted(result)))[0])

    def filtered(self):
        kept = self.kept()
        return [x for i, x in enumerate(self.result) if i in kept]

    def calc(self):
        filtered = self.filtered()
        if self.succ:
            result = len([x for x in filtered if self.succ(x)])
            if self.fail:
                result -= len([x for x in filtered if self.fail(x)])
            return result
        else:
            return sum(filtered)


class Parser(object):
    @classmethod
    def tokenize(cls, formula):
        whitespace = parsy.regex(r'\s*')
        number = parsy.regex(r'(0|[1-9][0-9]*)([.][0-9]+)?([eE][+-]?[0-9]+)?').map(float)
        expression = (whitespace >> (
            (parsy.regex(r'\*{2}|[()*/%+-]') | number | parsy.regex(r'\[.*?\]') |
             parsy.regex(r'\!{2}|\!p|mt|ro|k[hl]|d[hl]|s[ad]|[{}dFfmkdrs!,<>=]')
            ) << whitespace)).many()
        return expression.parse(formula)

    @classmethod
    def parse(cls, tokens):
        @parsy.generate
        def group_failures():
            result = yield group_successes
            yield parsy.match_item('f')
            comparison.operator_required = False
            condition = yield comparison
            result.fail = condition
            return result

        @parsy.generate
        def group_successes():
            result = yield group_drop
            comparison.operator_required = True
            condition = yield comparison | parsy.success('')
            if not condition:
                return result
            result.succ = condition
            return result

        @parsy.generate
        def group_drop():
            result = yield group_keep
            modifier = yield parsy.test_item(lambda x: x in ['d', 'dh', 'dl'], 'd[h|l]') | parsy.success('')
            if not modifier:
                return result
            count = yield number
            result.drop = (count, modifier == 'dh')
            return result

        @parsy.generate
        def group_keep():
            result = yield group
            modifier = yield parsy.test_item(lambda x: x in ['k', 'kh', 'kl'], 'k[h|l]') | parsy.success('')
            if not modifier:
                return result
            count = yield number
            result.keep = (count, modifier == 'kl')
            return result

        @parsy.generate
        def group():
            yield parsy.match_item('{')
            #result = yield group_simple
            result = yield expression_additive
            result = Group([result])
            while True:
                end = yield parsy.match_item('}') | parsy.success('')
                if end:
                    break
                yield parsy.match_item(',')
                #other = yield group_simple
                other = yield expression_additive
                result.items.append(other)
            return result

        @parsy.generate
        def expression_additive():
            result = yield expression_multiplicative
            sign = parsy.match_item('+') | parsy.match_item('-')
            while True:
                operation = yield sign | parsy.success('')
                if not operation:
                    break
                operand = yield expression_multiplicative
                if operation == '+':
                    result = Operation(operation, lambda x, y: x + y, result, operand)
                elif operation == '-':
                    result = Operation(operation, lambda x, y: x - y, result, operand)
            return result

        @parsy.generate
        def expression_multiplicative():
            result = yield expression_exponential
            operator = parsy.match_item('*') | parsy.match_item('/') | parsy.match_item('%')
            while True:
                operation = yield operator | parsy.success('')
                if not operation:
                    break
                operand = yield expression_exponential
                if operation == '*':
                    result = Operation(operation, lambda x, y: x * y, result, operand)
                elif operation == '/':
                    result = Operation(operation, lambda x, y: x / y, result, operand)
                elif operation == '%':
                    result = Operation(operation, lambda x, y: x % y, result, operand)
            return result

        @parsy.generate
        def expression_exponential():
            result = yield expression_simple
            operator = parsy.match_item('**')
            while True:
                operation = yield operator | parsy.success('')
                if not operation:
                    break
                operand = yield expression_simple
                if operation == '**':
                    result = Operation(operation, lambda x, y: x ** y, result, operand)
            return result

        @parsy.generate
        def roll():
            result = yield failures | successes
            label = parsy.test_item(lambda x: x.startswith('[') and x.endswith(']'), '[LABEL]')
            label = yield label | parsy.success(None)
            if label:
                label = label[1:-1]
            result.label = label
            return result

        @parsy.generate
        def failures():
            result = yield successes
            yield parsy.match_item('f')
            comparison.operator_required = False
            condition = yield comparison
            result.fail = condition
            return result

        @parsy.generate
        def successes():
            result = yield drop
            comparison.operator_required = True
            condition = yield comparison | parsy.success('')
            if not condition:
                return result
            result.succ = condition
            return result

        @parsy.generate
        def drop():
            result = yield keep
            modifier = yield parsy.test_item(lambda x: x in ['d', 'dh', 'dl'], 'd[h|l]') | parsy.success('')
            if not modifier:
                return result
            count = yield number
            result.drop = (count, modifier == 'dh')
            return result

        @parsy.generate
        def keep():
            result = yield modifiers
            modifier = yield parsy.test_item(lambda x: x in ['k', 'kh', 'kl'], 'k[h|l]') | parsy.success('')
            if not modifier:
                return result
            count = yield number
            result.keep = (count, modifier == 'kl')
            return result

        @parsy.generate
        def modifiers():
            result, dice = yield basic_roll
            modifications = {}
            while True:
                prio, modify = yield exploding | compounding | penetrating | reroll | sort | parsy.success((None, None))
                if not prio:
                    break
                if not prio in modifications:
                    modifications[prio] = []
                modifications[prio].append(modify)
            for prio, m in sorted(modifications.items()):
                if prio == 4: # reroll
                    changed = True
                    while changed:
                        for modify in m:
                            new_result, dice = modify(result, dice)
                            changed = new_result != result
                            result = new_result
                            if changed:
                                break
                else:
                    for modify in m:
                        result, dice = modify(result, dice)
            return Roll(result)

        @parsy.generate
        def sort():
            modifier = yield parsy.test_item(lambda x: x in ['s', 'sa', 'sd'], 's[a|d]')
            def modify(result, dice):
                result = sorted(result, reverse=(modifier == 'sd'))
                return result, dice
            return 5, modify

        @parsy.generate
        def reroll():
            modifier = yield parsy.test_item(lambda x: x in ['r', 'ro'], 'r[o]')
            comparison.operator_required = False
            condition = yield comparison
            def modify(result, dice):
                result = result[:]
                i = 0
                while i < len(result):
                    cnt = 0
                    while condition(result[i], dice[0]):
                        result[i] = random.choice(dice)
                        if modifier == 'ro':
                            break
                        cnt += 1
                        if cnt > ROLL_LIMIT:
                            raise RuntimeError('Roll limit reached!')
                    i += 1
                return result, dice
            return 4, modify

        @parsy.generate
        def penetrating():
            yield parsy.match_item('!p')
            comparison.operator_required = False
            condition = yield comparison
            def modify(result, dice):
                result = result[:]
                i = 0
                while i < len(result):
                    sub = [result[i]]
                    cnt = 0
                    while condition(sub[-1], dice[-1]):
                        sub.append(random.choice(dice))
                        cnt += 1
                        if cnt > ROLL_LIMIT:
                            raise RuntimeError('Roll limit reached!')
                    result[i+1 : i+1] = [x - 1 for x in sub][1:]
                    i += len(sub)
                return result, dice
            return 3, modify

        @parsy.generate
        def compounding():
            yield parsy.match_item('!!')
            comparison.operator_required = False
            condition = yield comparison
            def modify(result, dice):
                result = result[:]
                i = 0
                while i < len(result):
                    sub = [result[i]]
                    cnt = 0
                    while condition(sub[-1], dice[-1]):
                        sub.append(random.choice(dice))
                        cnt += 1
                        if cnt > ROLL_LIMIT:
                            raise RuntimeError('Roll limit reached!')
                    result[i] = sum(sub)
                    i += 1
                return result, dice
            return 2, modify

        @parsy.generate
        def exploding():
            yield parsy.match_item('!')
            comparison.operator_required = False
            condition = yield comparison
            def modify(result, dice):
                result = result[:]
                i = 0
                while i < len(result):
                    sub = [result[i]]
                    cnt = 0
                    while condition(sub[-1], dice[-1]):
                        sub.append(random.choice(dice))
                        cnt += 1
                        if cnt > ROLL_LIMIT:
                            raise RuntimeError('Roll limit reached!')
                    result[i: i+1] = sub
                    i += len(sub)
                return result, dice
            return 1, modify

        @parsy.generate
        def comparison():
            operator = parsy.match_item('<') | parsy.match_item('>') | parsy.match_item('=')
            if comparison.operator_required:
                operation = yield operator
            else:
                operation = yield operator | parsy.success('')
            if operation:
                operand = yield number
            else:
                operand = yield number | parsy.success('')
            def condition(value, default=None):
                if operation == '<':
                    return value <= operand
                elif operation == '>':
                    return value >= operand
                elif operation == '=':
                    return value == operand
                elif operand:
                    return value == operand
                else:
                    return value == default
            return condition

        @parsy.generate
        def basic_roll():
            count = yield computed_simple
            yield parsy.match_item('d')
            sides = yield computed_simple | parsy.match_item('F')
            dice = [-1, 0, 1] if sides == 'F' else [x + 1 for x in range(round(sides))]
            result = [random.choice(dice) for _ in range(round(count))]
            return result, dice

        @parsy.generate
        def computed_additive():
            result = yield computed_multiplicative
            sign = parsy.match_item('+') | parsy.match_item('-')
            while True:
                operation = yield sign | parsy.success('')
                if not operation:
                    break
                operand = yield computed_multiplicative
                if operation == '+':
                    result += operand
                elif operation == '-':
                    result -= operand
            return result

        @parsy.generate
        def computed_multiplicative():
            result = yield computed_exponential
            operator = parsy.match_item('*') | parsy.match_item('/') | parsy.match_item('%')
            while True:
                operation = yield operator | parsy.success('')
                if not operation:
                    break
                operand = yield computed_exponential
                if operation == '*':
                    result *= operand
                elif operation == '/':
                    result /= operand
                elif operation == '%':
                    result %= operand
            return result

        @parsy.generate
        def computed_exponential():
            result = yield computed_simple
            operator = parsy.match_item('**')
            while True:
                operation = yield operator | parsy.success('')
                if not operation:
                    break
                operand = yield computed_simple
                if operation == '**':
                    result **= operand
            return result

        @parsy.generate
        def number():
            sign = yield parsy.match_item('+') | parsy.match_item('-') | parsy.success('+')
            value = yield parsy.test_item(lambda x: isinstance(x, float), 'number')
            return num(value if sign == '+' else -value)

        computed_simple = (parsy.match_item('(') >> computed_additive << parsy.match_item(')')) | number
        expression_simple = roll | (parsy.match_item('(') >> expression_additive << parsy.match_item(')')) | number
        group_simple = group_failures | group_successes | expression_additive
        return group_simple.parse(tokens)


class Roll20Error(Exception):
    pass


class Roll20(object):
    @classmethod
    def execute(cls, formula):
        try:
            tokens = Parser.tokenize(formula)
            result = Parser.parse(tokens)
            try:
                calculated = result.calc()
            except AttributeError:
                calculated = num(result)
            return '{0} = __**{1}**__'.format(str(result), calculated)
        except (parsy.ParseError, RuntimeError) as e:
            raise Roll20Error(str(e))