# Filename: simulation.py
# Author: Raiyan Rahman
# Date: Jan 18th, 2023
# Description: Selecting an outcome according to a given probability
# distribution is a common task faced by simulation programs.
# The following program demonstrates one way of accomplishing this task.

import random

# A fair coin has 50%-50% chance of heads or tails.
FAIR_COIN_PROBDIST = [1/2, 1/2]
# A fair die has 1/6 chance of landing on 1, 2, 3, 4, 5, and 6 each.
FAIR_DIE_PROBDIST = [1/6, 1/6, 1/6, 1/6, 1/6, 1/6]

def distro_select(prob_distribution: list) -> int:
    """
    Return the outcome of running a simulation based on the given
    probability distribution. Each index in the prob_distribution list is the
    probability of a separate possible result.
    :param prob_distribution: the probability distribution of each outcome
    :return: result after running simulation
    """
    # Get the list of possible outcomes.
    possible_outcomes = [i for i in range(len(prob_distribution))]
    # Select result from the possible outcomes given prob distribution.
    results = random.choices(population=possible_outcomes,
                             weights=prob_distribution,
                             k=1)
    # Return the result.
    return results[0]


if __name__ == '__main__':
    num_trials = 100 # Set how many results you want here.

    # Fair coin case.
    print(f'Tossing a fair coin {num_trials} times:')
    # Run the trials.
    for trial in range(num_trials):
        # End parameter avoids newline.
        print(distro_select(prob_distribution=FAIR_COIN_PROBDIST), end=' ')
    print('\n') # Extra newline.

    # Fair die case.
    print(f'Throwing a fair die {num_trials} times:')
    # Run the trials.
    for trial in range(num_trials):
        # End parameter avoids newline.
        print(distro_select(prob_distribution=FAIR_DIE_PROBDIST), end=' ')
    print()