horse_sim.py

from multiprocessing import Pool
from random import random

import matplotlib.pyplot as plt
import pandas as pd
from tqdm import tqdm

TIER_CNT = 4
WORKER_CNT = 8
SIM_CNT = 1_000_000

class Sim:
    def __init__(self):
        self.stock = [0] * TIER_CNT
        self.browns = 0

    def consume(self, tier: int):
        if self.stock[tier] > 0: # we have one in stock
            self.stock[tier] -= 1
        elif tier == 0: # it's brown, so buy one
            self.browns += 1
        else: # breed horses until one of the right tier is available
            while True:
                self.consume(tier - 1)
                self.consume(tier - 1)
                roll = random()
                if roll < 0.2: # success
                    return
                elif roll < 0.5: # tier deteriorated
                    self.store(tier - 2)
                else: # tier didn't change
                    self.store(tier - 1)

    def store(self, tier: int):
        self.stock[max(tier, 0)] += 1

def run_sim(_: int) -> int:
    sim = Sim()
    sim.consume(TIER_CNT - 1)
    return sim.browns

def main():
    # compute
    worker_pool = Pool(WORKER_CNT)
    result_iter = worker_pool.imap_unordered(run_sim, range(SIM_CNT), chunksize=10000)
    worker_pool.close()
    datapoint_list = list(tqdm(result_iter, total=SIM_CNT, leave=True, unit='iter'))
    
    # analyze
    dataset = pd.DataFrame(datapoint_list)
    print(dataset.describe())
    dataset.plot.hist(bins=256)
    plt.show()

if __name__ == '__main__':
    main()