Vermintide_dice_prob

#!/usr/bin/env
# -*- coding: utf-8 -*-
# © Konstantin Soin <[email protected]>

import pylab
import random
from decimal import Decimal

class Cube(object):

    def __init__(self, type='default'):
        self.side_one = True
        self.side_two = True
        if type == 'default' or type == 2:
            self.type = 'default'
            self.side_three = False
            self.side_four = False
            self.side_five = False
            self.side_six = False
        elif type == 'fair' or type == 3:
            self.type = 'fair'
            self.side_three = True
            self.side_four = False
            self.side_five = False
            self.side_six = False
        elif type == 'four' or type == 4:
            self.type = 'four'
            self.side_three = True
            self.side_four = True
            self.side_five = False
            self.side_six = False
        elif type == 'win' or type == 6:
            self.type = 'win'
            self.side_three = True
            self.side_four = True
            self.side_five = True
            self.side_six = True

    def change_type(self, type):
        """Changes default cube type to selected: 3 or 4 or 6"""
        if type == 'default' or type == 2:
            self.type = 'default'
            self.side_three = False
            self.side_four = False
            self.side_five = False
            self.side_six = False
        elif type == 'fair' or type == 3:
            self.type = 'fair'
            self.side_three = True
            self.side_four = False
            self.side_five = False
            self.side_six = False
        elif type == 'four' or type == 4:
            self.type = 'four'
            self.side_three = True
            self.side_four = True
            self.side_five = False
            self.side_six = False
        elif type == 'win' or type == 6:
            self.type = 'win'
            self.side_three = True
            self.side_four = True
            self.side_five = True
            self.side_six = True

    def roll(self):
        """Randomly roll the dice"""
        roll = random.ranint(1, 6)
        if roll == 1 or roll == 2:
            return True
        elif roll == 3:
            return self.side_three
        elif roll == 4:
            return self.side_four
        elif roll == 5:
            return self.side_five
        elif roll == 6:
            return self.side_six

    def get_type(self):
        print self.type

    def get_sides(self):
        return [self.side_one, self.side_two, self.side_three, self.side_four, self.side_five, self.side_six]

class Cube_pull(object):
    def __init__(self, loot_dices=0, books=0, grimoires=0):
        self.cubes = []
        self.cc = 7
        for i in xrange(loot_dices):
            self.cubes.append(Cube('fair'))
            self.cc -= 1
        for i in xrange(books):
            self.cubes.append(Cube('four'))
            self.cc -= 1
        for i in xrange(grimoires):
            self.cubes.append(Cube('win'))
        for i in xrange(self.cc):
            self.cubes.append(Cube())

    def get_vars(self):
        vars = []
        for side_0 in self.cubes[0].get_sides():
            for side_1 in self.cubes[1].get_sides():
                for side_2 in self.cubes[2].get_sides():
                    for side_3 in self.cubes[3].get_sides():
                        for side_4 in self.cubes[4].get_sides():
                            for side_5 in self.cubes[5].get_sides():
                                for side_6 in self.cubes[6].get_sides():
                                    vars.append([side_0, side_1, side_2, side_3, side_4, side_5, side_6])
        return vars

    def get_prob(self):
        count_full = 0
        count_zero = 0
        count_one = 0
        count_two = 0
        count_three = 0
        count_four = 0
        count_five = 0
        count_six = 0
        count_seven = 0

        for var in self.get_vars():
            count = var.count(True)
            count_full += 1
            if count == 0:
                count_zero += 1
            elif count == 1:
                count_one += 1
            elif count == 2:
                count_two += 1
            elif count == 3:
                count_three += 1
            elif count == 4:
                count_four += 1
            elif count == 5:
                count_five += 1
            elif count == 6:
                count_six += 1
            elif count == 7:
                count_seven += 1
        count_full = Decimal(count_full)
        return [count_zero/count_full, count_one/count_full, count_two/count_full, count_three/count_full, count_four/count_full, count_five/count_full, count_six/count_full, count_seven/count_full]


if __name__ == "__main__":
    pylab.figure(1)
    my_pull = Cube_pull(0, 2, 1) # loot dices, books, grimoireses
    my_probs = my_pull.get_prob()
    for p in my_probs:
        print '%d win dices = %05r %%' % (my_probs.index(p), round(p*100, 2))
    pylab.plot(range(8), my_probs)
    pylab.show()