import math
import matplotlib.pyplot as plt
import numpy as np
import random
import scipy.stats as stats

die_e = 3.5 # expected value of an ordinary six-sided die
die_var = 17.5/6 # variance of an ordinary six-sided die
die_sigma = math.sqrt(die_var)

def create_statistics(repetitions, casts):
    clt_sum_frequencies = np.zeros(6 * casts + 1, dtype = int)
    total_sum = 0
    for repetition in range(0, repetitions):
        sum = 0
        for cast in range(0, casts):
            die_number = random.randint(1, 6)
            sum = sum + die_number
        clt_sum_frequencies[sum] += 1
        total_sum += sum
    return clt_sum_frequencies / repetitions

def plot_clt_chart(sum_frequencies, shift, scale, casts, ax, title):
    sums = np.linspace((0 - shift) / scale, (6 * casts - shift) / scale, len(sum_frequencies))
    width = sums[1] - sums[0]
    sum_densities = sum_frequencies / width
    mu = (casts * die_e - shift) / scale
    sigma = die_sigma * math.sqrt(casts) / scale
    normal_densities = [stats.norm.pdf(x, mu, sigma) for x in sums]
    plot_chart(sums, sum_densities, normal_densities, width, ax, title)

def plot_chart(x_values, y_values, y_comparison_values, width, ax, title):
    x_values_odd = x_values[np.mod(np.arange(x_values.size), 2) != 0]
    y_values_odd = y_values[np.mod(np.arange(x_values.size), 2) != 0]
    x_values_even = x_values[np.mod(np.arange(y_values.size), 2) == 0]
    y_values_even = y_values[np.mod(np.arange(y_values.size), 2) == 0]
    ax.bar(x_values_odd, y_values_odd, width=width, align='center', color='blue')
    ax.bar(x_values_even, y_values_even, width=width, align='center', color='palegreen')
    ax.plot(x_values, y_comparison_values, color='r')
    ax.set_title(title)

fig, ax = plt.subplots(2, sharex=True, sharey=True, figsize=(5, 10))
plt.xlim(-0.5, 0.5)
row = 0
for casts in [100, 1000]:
    for repetitions in [10000]:
        title = f"{repetitions} times averages of {casts} random numbers"
        clt_sum_frequencies = create_statistics(repetitions, casts)
        plot_clt_chart(clt_sum_frequencies, die_e * casts, casts, casts, ax[row], title)
    row += 1
plt.show()