shadow_combat.py

"""Shadow run attack simulator.

Usage:
  shadow_combat.py [--debug] [--limit LIMIT | -6 | --rule_of_six] [--threshold THRESHOLD | --opposed_pool OPPOSED [--opposed_limit OPPOSED_LIMIT]] [--dv DV --stun [--soak SOAK] [--armor ARMOR [--ap AP]]] [--contact] [--once | -o | [--iterations ITERATIONS] [-D | --distribution]] [--multi ATTACKS] [--min DAMAGE] ATTACK_POOL
  shadow_combat.py [--debug] [--contact] [-6 | --rule_of_six] --threshold THRESHOLD ATTACK_STRING DAMAGE_STRING SOAK_STRING
  shadow_combat.py [--debug] [--contact] [-D | --distribution] [-6 | --rule_of_six] --threshold THRESHOLD ATTACK_STRING DAMAGE_STRING SOAK_STRING ITERATIONS
  shadow_combat.py [--debug] [--contact] [-6 | --rule_of_six] ATTACK_STRING DAMAGE_STRING DEFENSE_STRING SOAK_STRING
  shadow_combat.py [--debug] [--contact] [-D | --distribution] [-6 | --rule_of_six] ATTACK_STRING DAMAGE_STRING DEFENSE_STRING SOAK_STRING ITERATIONS
  shadow_combat.py (-h | --help)
  shadow_combat.py (-v | --version)

Options:
  -6 --rule_of_six              rule of six applies
  --limit LIMIT                 attack dice limit
  --threshold THRESHOLD         threshold for test
  --opposed_pool OPPOSED        opposed dice pool
  --opposed_limit OPPOSED_LIMIT opposed dice limit
  --dv DV                       the damage value of the attack
  --soak SOAK                   defense base soak dice (does not include armor)
  --armor ARMOR                 armor dice
  --ap AP                       armor piercing [default: 0]
  --stun                        dv is stun damage, not physical
  --contact                     if this is a contact only attack [default: False]
  --multi ATTACKS               split attack over ATTACKS [default: 1]
  -o --once                     simulate a single action only
  -D --distribution             print the distribution of attacks
  --iterations ITERATIONS       total iterations to run through [default: 10000]
  -h --help                     show this
  -v --version                  print version
  --debug                       turn on debug printing
  --min DAMAGE                  minimum_damage goal

ATTACK_STRING format is 18[7] meaning 18 dice in the pool, with a limit of 7
DAMAGE_STRING format is 7Pv-4 meaning 7 physical damage, with -4 AP
DEFENSE_STRING format is same as ATTACK_STRING... without a limit it'd just be 18
SOAK_STRING format is 12v3 meaning 12 armor, 3 soak dice or just v3 for the case where there is no armor involved

"""

from functools import total_ordering
from itertools import groupby, chain
import re
from numbers import Number
from sys import argv, stderr, stdout
from warnings import warn
    
from docopt import docopt

try:
    import random
    sys_random = random.SystemRandom()
    choice, randint = sys_random.choice, sys_random.randint
except:
    warn('Could not use SystemRandom; using default random instead', RuntimeWarning)
    from random import choice, randint

class InputError(Exception):
    """Exception raised for errors in input

    Attributes:
        expr -- input expression in which the error occurred
        msg  -- explanation of the error
    """

    def __init__(self, expr, mesg):
        self.expr = expr
        self.msg = mesg


@total_ordering
class Outcome(object):
    'Represents the outcome of one or more attacks'

    __slots__ = ('hits', 'phys', 'stun', 'misses')

    SINGULAR = {
        'misses': 'miss',
        'hits': 'hit',
        'phys': 'phys',
        'stun': 'stun'
    }

    def __init__(self, misses=0, hits=None, phys=0, stun=0, copy=None):
        if copy is None:
            assert isinstance(misses, int) and isinstance(phys, int) and isinstance(stun, int)
            self.misses = misses
            self.phys = phys
            self.stun = stun
            # if hits is None, set it to an implied value based on damage
            if hits is None:
                self.hits = 1 if (self.phys > 0 or self.stun > 0) else 0
            else:
                self.hits = hits
            assert isinstance(self.hits, int)
        else:
            assert misses == 0 and hits is None and phys == 0 and stun == 0
            assert isinstance(copy, Outcome)
            self.misses = copy.misses
            self.hits = copy.hits
            self.phys = copy.phys
            self.stun = copy.stun

    def __iadd__(self, other):
        self.misses += other.misses
        self.hits += other.hits
        self.phys += other.phys
        self.stun += other.stun
        return self

    def __isub__(self, other):
        self.misses -= other.misses
        self.hits -= other.hits
        self.phys -= other.phys
        self.stun -= other.stun
        return self

    def __imul__(self, other):
        assert isinstance(other, Number)
        self.misses *= other
        self.hits *= other
        self.phys *= other
        self.stun *= other
        return self

    def __itruediv__(self, other):
        assert isinstance(other, Number)
        self.misses /= other
        self.hits /= other
        self.phys /= other
        self.stun /= other
        return self

    def __add__(self, other):
        return Outcome(copy=self).__iadd__(other)

    def __sub__(self, other):
        return Outcome(copy=self).__isub__(other)

    def __mul__(self, other):
        return Outcome(copy=self).__imul__(other)

    def __truediv__(self, other):
        return Outcome(copy=self).__itruediv__(other)

    def __radd__(self, other):
        assert other == 0
        return Outcome(copy=self)

    def __rsub__(self, other):
        assert other == 0
        return Outcome() - Outcome(copy=self)

    def __rmul__(self, other):
        return self.__mul__(other)

    def __rtruediv__(self, other):
        return self.__truediv__(other)

    @staticmethod
    def singular(attribute):
        return Outcome.SINGULAR.get(attribute, attribute)

    def __str__(self):
        if self.hits > 0 and self.stun == 0 and self.phys == 0:
            return '{:>5} {:>4},  no damage'.format(self.hits, 'hit' if self.hits == 1 else 'hits')
        answer = ', '.join('{:>5} {:>8}'.format(*x) for x in ((getattr(self, attr), (self.singular(attr) if getattr(self, attr) == 1 else attr)) for attr in self.__slots__) if x[1] == 'miss' or x[0] > 0)
        return '{:22}'.format(answer)

    def __repr__(self):
        return 'Outcome({}, {}, {}, {})'.format(self.misses, self.hits, self.phys, self.stun)

    def __lt__(self, other):
        if self.phys < other.phys:
            return True
        if self.phys == other.phys:
            if self.stun < other.stun:
                return True
            if self.stun == other.stun:
                if self.hits < other.hits:
                    return True
                if self.hits == other.hits:
                    return self.misses > other.misses # more misses means overall you did worse
        return False

    def __eq__(self, other):
        return self.misses == other.misses and self.hits == other.hits and self.phys == other.phys and self.stun == other.stun


class Pool(object):
    OUTCOMES = (False, False, True) # we could do False, False, False, False, True, True... but that is needlessly inefficient
    def __init__(self, dice, limit=None, rule_of_six=False):
        if rule_of_six:
            assert limit is None
        self.dice = int(dice)
        self.limit = int(limit) if limit else None
        self.rule_of_six = rule_of_six

    def hit(self):
        return choice(self.OUTCOMES)

    def hit_generator(self):
        if self.rule_of_six:
            dice = self.dice
            while dice > 0:
                roll = randint(1, 6)
                yield roll > 4
                if roll != 6:
                    dice -= 1

        else:
            for _ in range(self.dice):
                yield self.hit()

    def hits(self):
        hits = sum(1 for hit in self.hit_generator() if hit)
        return hits if self.limit is None else min(hits, self.limit)

    def split(self, parts=1):
        base = self.dice // parts
        more_than_base = (base + 1 for _ in range(self.dice % parts))
        rest = (base for _ in range(parts - (self.dice % parts)))
        return (Pool(dice, self.limit, self.rule_of_six) for dice in chain(more_than_base, rest))

    def __str__(self):
        return '{}[{}]'.format(self.dice, self.limit) if self.limit else '{}{}'.format(self.dice, '*' if self.rule_of_six else '')

    def __repr__(self):
        return 'Pool(dice={}, limit={}, rule_of_six={})'.format(self.dice, self.limit, self.rule_of_six)


class FixedThreshold(object):
    def __init__(self, threshold):
        self.threshold = int(threshold)

    def hits(self):
        return self.threshold

    def __str__(self):
        return '({})'.format(self.threshold)

    def __repr__(self):
        return 'FixedThreshold({})'.format(self.threshold)


class Test(object):
    def __init__(self, pool, threshold=None):
        self.pool = pool
        self.threshold = threshold

    def test(self):
        hits = self.pool.hits()
        return hits if self.threshold is None else hits - self.threshold.hits()

    def net_hits(self):
        return max(0, self.test())

    def __str__(self):
        return '{} vs {}'.format(self.pool, self.threshold)

    def __repr__(self):
        return 'Test(pool={}, threshold={})'.format(repr(self.pool), repr(self.threshold))


class Attack(object):
    def __init__(self, test, dv=None, ap=0, stun=False, soak=None, armor=None, contact=False):
        self.test = test
        self.dv = int(dv) if dv else None
        self.stun = stun
        self.soak = int(soak) if soak else None
        self.armor = int(armor) if armor else None
        self.ap = int(ap) if ap else 0
        self.contact = contact

    def damage(self):
        hits = self.test.net_hits()
        if (hits < 0) or (hits == 0 and not self.contact):
            return Outcome(misses=1)

        # if dv is not set, then we're just counting hits vs. not hits
        if self.dv is None:
            return Outcome(hits=1)

        modified_dv = self.dv + hits

        if self.soak is None:
            return Outcome(stun=modified_dv) if self.stun else Outcome(phys=modified_dv)

        stun = self.stun

        soak = self.soak

        if self.armor is not None:
            armor_value = self.armor + self.ap
            if armor_value > 0:
                if armor_value >= modified_dv:
                    stun = True
                soak += armor_value

        soak_pool = Pool(soak)
        damage = max(0, modified_dv - soak_pool.hits())
        return Outcome(hits=1, stun=damage) if stun else Outcome(hits=1, phys=damage)

    def __str__(self):
        ap_val = self.ap if self.ap < 0 else ('+' + self.ap if self.ap > 0 else '')
        base = '{}{} {}{}{}'.format(self.test,
                                    ' contact' if self.contact else '',
                                    self.dv,
                                    'S' if self.stun else 'P',
                                    ap_val)
        return base + 'soak: {}v{}'.format(self.armor, self.soak) if self.armor or self.soak else base

    def __repr__(self):
        return 'Attack(test={}, dv={}, ap={}, stun={}, soak={}, armor={}, contact={})'.format(repr(self.test), repr(self.dv), repr(self.ap), repr(self.stun), repr(self.soak), repr(self.armor), repr(self.contact))


class Simulation(object):
    def __init__(self, attacks, iterations=10000, debug=False):
        self.attacks = tuple(attacks)
        if debug:
            stderr.write('Performing {} iterations of [{}]\n'.format(iterations, ', '.join(map(str, self.attacks))))
        self.iterations = int(iterations) if iterations else 10000
        self.outcome = None

    def results(self):
        for _ in range(self.iterations):
            d = sum(attack.damage() for attack in self.attacks)
            assert d.misses != 0 or d.hits != 0
            yield d

    def distribution(self):
        pred = lambda x: (x.phys, x.stun)
        results = sorted(self.results(), reverse=True)
        for damage, g in groupby(results):
            total = sum(1 for _ in g)
            yield damage, total
            if self.outcome is None:
                self.outcome = damage * total
            else:
                self.outcome += damage * total

    def mean(self):
        return self.outcome * 1.0 / self.iterations

ATTACK_RE = re.compile(r'^(\d+)(?:\[(\d+)\])?$')
DEFENSE_RE = ATTACK_RE
DAMAGE_RE = re.compile(r'^(\d+)(P|S)(?:v((?:\+|-)\d+))?$')
SOAK_RE = re.compile(r'^(\d+)?v(\d+)$')

def get_attack_pool(args):
    if args['ATTACK_POOL']:
        return Pool(args['ATTACK_POOL'], args['--limit'], args['--rule_of_six'])
    elif args['ATTACK_STRING']:
        matcher = ATTACK_RE.match(args['ATTACK_STRING'])
        if matcher is None:
            raise InputError(args['ATTACK_STRING'], 'Invalid Attack String')
        limit = matcher.group(2)
        return Pool(matcher.group(1), limit or None, args['--rule_of_six'])
    else:
        raise InputError(' '.join(args), 'No attacker parameters found')

def get_threshold(args):
    if args['--threshold']:
        return FixedThreshold(args['--threshold'])
    if args['--opposed_pool']:
        return Pool(args['--opposed_pool'], args['--opposed_limit'])
    if args['DEFENSE_STRING']:
        matches = DEFENSE_RE.match(args['DEFENSE_STRING'])
        if matches is None:
            raise InputError(args['DEFENSE_STRING'], 'Invalid Defense String')
        limit = matches.group(2)
        return Pool(matches.group(1), limit or None)
    else:
        return None

def parse_ap_string(ap_string):
    if ap_string[0] == '+':
        return int(ap_string[1:])
    return -int((ap_string[1:] if ap_string[0] == '-' else ap_string))

def get_tests(args):
    attack_pool = get_attack_pool(args)
    threshold = get_threshold(args)
    multi_attacks = int(args['--multi']) if args['--multi'] else 1
    if args['--debug']:
        stderr.write('Splitting in to {} attacks\n'.format(multi_attacks))
    return (Test(pool, threshold) for pool in attack_pool.split(multi_attacks))

def get_attacks(tests, args):
    attack_args = {}
    if args['DAMAGE_STRING'] and args['SOAK_STRING']:
        matches = DAMAGE_RE.match(args['DAMAGE_STRING'])
        if matches is None:
            raise InputError(args['DAMAGE_STRING'], 'Invalid Damage String')
        assert matches.group(2) in ('S', 'P')
        attack_args['stun'] = matches.group(2) == 'S'
        ap_string = matches.group(3)
        attack_args['ap'] = None if ap_string is None else parse_ap_string(ap_string)
        attack_args['dv'] = int(matches.group(1))
        matches = SOAK_RE.match(args['SOAK_STRING'])
        if matches is None:
            raise InputError(args['SOAK_STRING'], 'Invalid Soak String')
        attack_args['armor'] = None if matches.group(1) is None else int(matches.group(1))
        attack_args['soak'] = int(matches.group(2))
    else:
        attack_args['dv'] = args['--dv']
        attack_args['ap'] = parse_ap_string(args['--ap'])
        attack_args['stun'] = args['--stun']
        attack_args['soak'] = args['--soak']
        attack_args['armor'] = args['--armor']
    attack_args['contact'] = args['--contact']
    for test in tests:
        yield Attack(test, **attack_args)

def main(**args):
    if args['--debug']:
        stderr.write('Args: {}\n'.format(args))

    tests = get_tests(args)
    attacks = get_attacks(tests, args)

    if args['--once']:
        for attack in attacks:
            print(attack.damage())
        return

    simulation = Simulation(attacks, args['--iterations'] or args['ITERATIONS'], args['--debug'])

    distribution = simulation.distribution()
    min_damage = args['--min']
    if min_damage:
        min_damage = int(min_damage)
        min_outcomes = 0
    for outcome, total in distribution:
        if min_damage and outcome.phys >= min_damage:
            min_outcomes += total
        if args['--distribution']:
            print('{} for {} or {:>5.2f}%'.format(outcome, total, total * 100.0 / simulation.iterations))

    if min_damage:
        print('{:>5.2f}% above {}P'.format(min_outcomes * 100.0 / simulation.iterations, min_damage))
    print('mean: {}'.format(simulation.mean()))

if __name__ == '__main__':
    arguments = docopt(__doc__, version='Rolls 0.0')
    main(**arguments)