risk-calculator.py

#!/usr/bin/env python

from __future__ import division
import sys

def cross(*args):
    if not args:
        yield []
        return
    for i in args[0]:
        for n in cross(*(args[1:])):
            yield tuple([i] + list(n))

# http://linux.derkeiler.com/Mailing-Lists/Fedora/2005-11/4432.html
def combinations_r( list_of_items, number_of_picks ):
    k = number_of_picks
    n = len(list_of_items)
    if k==0 or n==0:
        return
    b = [0]*k # make list [0,0,0,....,0] - of length k
    bdone = [0]*k
    while True:
        choice = [list_of_items[i] for i in b]
        yield tuple(choice)
        # Modify b in place (an addition operation in base-n)
        for i in range(k):
            b[i] += 1
            if b[i] < n:
                break # break our of for loop, no need to carry
            b[i] = 0 # carry
        if b == bdone:
            break # break out of while loop, no more choices left
    return

RANGE = range(1,7)

DIE_PROB = 1/max(RANGE)

SINGLE_DIE = list((x,) for x in RANGE)
TWO_DICE = list(combinations_r(RANGE, 2))
THREE_DICE = list(combinations_r(RANGE, 3))

dice_rolls = {
    1: SINGLE_DIE,
    2: TWO_DICE,
    3: THREE_DICE,
}

def risk_wins(attack, defense):
    at_victory = 0
    de_victory = 0
    for at,de in zip( sorted(attack, lambda x, y: y-x),
                      sorted(defense, lambda x, y: y-x) ):
        if at > de:
            at_victory += 1
        else:
            de_victory += 1
    return at_victory, de_victory


def make_probability(nattack, ndefense):
    attacker = dice_rolls[nattack]
    defender = dice_rolls[ndefense]

    probs = {
        (2, 0): 0,
        (1, 0): 0,
        (1, 1): 0,
        (0, 1): 0,
        (0, 2): 0,
    }
    values = list(cross(attacker, defender))
    for at_roll, de_roll in values:
        at, de = risk_wins(at_roll, de_roll)
        probs[(at,de)] += 1

    l = len(values)
    for k in probs:
        probs[k] /= l

    return probs

def calculate_round(nattack, ndefense):
    probs = make_probability(nattack, ndefense)
    diffs = {}
    for k, v in probs.items():
        at, de = k
        diffs[(-de, -at)] = v
    return diffs

def decision_sum(decision):
    ats, des = 0, 0
    cum_prob = 1.0
    for (at, de), prob in decision:
        ats += at
        des += de
        cum_prob *= prob
    return ats, des, cum_prob

def make_decisions(nattack, ndefense, iterations):
    decisions = []
    from pprint import pprint
    for i in range(iterations):
        if not decisions:
            decisions = [ [(o, p)] for o, p in calculate_round(min(nattack-1, 3), min(ndefense, 2)).items() if p != 0.0]
            continue
        for decision in decisions[:]:
            at, de, _ = decision_sum(decision)
            try:
                round = calculate_round(min(nattack+at-1, 3), min(ndefense+de, 2))
            except (KeyError, AssertionError):
                continue
            for outcome, prob in round.items():
                if prob == 0.0:
                    continue
                try:
                    decisions.remove(decision)
                except:
                    pass
                decisions.append(decision + [(outcome,prob)])
    return decisions

def calculate_outcomes(nattack, ndefense, iterations):
    decisions = make_decisions(nattack, ndefense, iterations)
    results = []
    for decision in decisions:
        ats, des, cum_prob = decision_sum(decision)
        if cum_prob == 0.0:
            continue
        results.append( (nattack+ats, ndefense+des, cum_prob) )
    assert str(sum(prob for na, nd, prob in results)) == "1.0"
    return results

def print_outcomes(results, n):
    winp = 0.0
    losep = 0.0
    draws = {}
    for nat, nde, prob in results:
        if nat == 1:
            losep += prob
        elif nde == 0:
            winp += prob
        else:
            if (nat, nde) not in draws:
                draws[(nat,nde)] = 0
            draws[(nat,nde)] += prob
    draws = [ (nat, nde, prob) for (nat, nde), prob in draws.items() ]

    print "Chance of victory in %d battles: %3.2f%%" % (n, winp*100)
    print "Chance of defeat (1 remaining) in %d battles: %3.2f%%" % (n, losep*100)
    for at, de, prob in draws:
        print "%3.2f%% chance of:" % (prob*100)
        print "\tAttacker remaining: %d" % at
        print "\tDefender remaining: %d" % de

def print_result(probs):
    for result, prob in probs.items():
        if prob == 0.0:
            continue
        at, de = result
        print "Attacker Loses: %d, Defender Loses: %d (probability %1.2f)" % (de, at, prob)

def testmain():
    for at, de in cross(range(1,4),range(1,3)):
        print "------------------------------------------------------------------------"
        print " Attacker: %d dice Defender: %d dice" % (at, de)
        print "------------------------------------------------------------------------"
        print_result(make_probability(at, de))

def main():
    at, de, n = int(sys.argv[1]), int(sys.argv[2]), int(sys.argv[3])
    print "Attacking Units: %d\t\tDefending Units: %d\t\t# of Battles: %d" % (at, de, n)
    print_outcomes(calculate_outcomes(at, de, n), n)

if __name__=='__main__':
    sys.exit(main())