import sys
import random
from math import exp
from matplotlib import pyplot as plt

"""
Statistical results of random distribution:
"""

#=============================
# 大域変数の定義
#=============================
# Number of persons
npersons = 50
# The initial money that the persons have, equal to the average money
value    = 50
# The number of maximum iteration
niter    = 500
# The cycle to update graph
nplotinterval = 100

# The number of x mesh for distribution function
nfx = 21


# Treat argments
argv = sys.argv
if len(argv) == 1:
    print("Usage: python randomtrade.py npersons value(average) n(maxiteration) n(plotinterval) n(distribution func)")
    print("   ex: python randomtrade.py 200 50 10000 100 21")
    exit()

if len(argv) >= 2:
    npersons = int(argv[1])
if len(argv) >= 3:
    value = int(argv[2])
if len(argv) >= 4:
    niter = int(argv[3])
if len(argv) >= 5:
    nplotinterval = int(argv[4])
if len(argv) >= 6:
    nfx = int(argv[5])


def main():
# 初期化
    xindex = range(npersons)  # list of indexes of the persons to be used for x axis 
    v = []               # list of money the persons have
    for i in range(npersons):
        v.append(value)

#=============================
# グラフの表示
#=============================
# theoretical distribution: f = A * exp(-bv)
# A and b are determined by the number of persons and the money average
    print("")
    print("Analyatical distribution:")
    xmax = value * 10.0
    nftx = 101
    xstep = xmax / (nftx - 1)
    b = 1.0 / value
    A = npersons * b
    print("  x range: 0 - %g" % (xmax))
    print("  x step: %g" % (xstep))
    print("  f(x) = %g * exp(-%g * x)" % (A, b))
    ftx = []
    fty = []
    for i in range(nftx+1):
        x = i * xstep
        ftx.append(x)
        fty.append(A * exp(-b * x))

# Prepare plot
    fig = plt.figure()

    ax1 = fig.add_subplot(3, 1, 1)
    ax2 = fig.add_subplot(3, 1, 2)
    ax3 = fig.add_subplot(3, 1, 3)
    plt.subplots_adjust(wspace = 0.2, hspace = 0.3)

# randomize
    for i in range(niter):
        v = []
        sum = 0.0
        for ip in range(npersons):
            val = int(random.random() * value)
            sum += val
            v.append(val)

        k = value * npersons / sum
        for ip in range(npersons):
            v[ip] *= k

        vs = sorted(v)

        if i == 0 or (i+1) % nplotinterval == 0 or i == niter - 1:
# calculate distribution function f(x)
            vmax = max(v)
            vstep = vmax / (nfx - 1)

            fx  = []
            fy  = []
            for ifx in range(nfx):
                vlim0 = ifx * vstep
                vlim1 = (ifx+1) * vstep
                fx.append(vlim0)
                fy.append(0.0)
                for ip1 in range(npersons):
                    if vlim0 < vs[ip1] <= vlim1:
                        fy[ifx] += 1.0 / vstep

            ax1.cla()
            ax2.cla()
            ax3.cla()

            ax1.bar(xindex, v, width = 1.5)
            ax1.plot([0, npersons], [value, value], linestyle = 'dashed', color = 'red', linewidth = 0.5)
            ax1.set_title('Random distribution (i = %d/%d, n=%d, avg=%d)' % (i+1, niter, npersons, value))
            ax1.set_xlabel("i")
            ax1.set_ylabel("value")
            ax1.set_xlim([0, npersons])

            ax2.bar(xindex, vs, width = 1.5)
            ax2.plot([0, npersons], [value, value], linestyle = 'dashed', color = 'red', linewidth = 0.5)
            ax2.set_xlim([0, npersons])
            ax2.set_xlabel("i")
            ax2.set_ylabel("value (sorted)")

            ax3.plot(fx, fy, label = 'distrib')
            ax3.plot(ftx, fty, label = 'exponential', color = "red", linewidth = 0.5, linestyle = 'dashed')
            ax3.set_xlim([0, vmax])
            ax3.set_ylim([0, max(fy) + 1])
            ax3.set_xlabel("value")
            ax3.set_ylabel("frequency")
            ax3.legend()

            plt.pause(0.1)

    print("Press ENTER to exit>>", end = '')
    input()


if __name__ == '__main__':
    main()