#!/usr/bin/env python3
import sys
import matplotlib.pyplot as plt
import numpy as np

# This script generates a graph visualizing the theoretical difference
# between the pi-blocking of the NJLP and FMLP for the input values
# of n and m.

# The n-th harmonic number is the sum(1/i) for i=[1,n]
def harmonic_between(n, m):
	sum = 0;
	for i in range(m+1, n+1):
		sum += float(1)/i
	return sum

# Output the arguments
if len(sys.argv) != 4:
	print("Usage: plot_blocking lg2mmax nmax lmax")

# lg2mmax controls which values for m we generate this
# graph for. For example, lg2mmax=4 means we generate
# graphs for m=1,2,4,8
lg2mmax = int(sys.argv[1])
# Creates a data point for every n in [1,nmax]
nmax = int(sys.argv[2])
# The largest critical section duration
lmax = float(sys.argv[3])

# The colors of the four graphs
colors = ['r', 'g', 'b', 'y']
# Tick marks for y axis
yticks = np.arange(lg2mmax)
ms = np.exp2(lg2mmax - 1 - yticks).astype(int)

fig = plt.figure()
ax = fig.add_subplot(projection='3d')

for c, k, m in zip(colors, yticks, ms):
	xs = [None] * nmax
	ys = [None] * nmax

	print(m)

	for n in range(1, nmax + 1):
        # The NJLP pi-blocking bound is ( 3m-1+m(H_n - H_m) )*Lmax
		njlp_blocking = (3*m - 1 + m*(harmonic_between(n, m)))*lmax
        # The FMLP blocking bound is ( n-1 )*Lmax
		fmlp_blocking = (n - 1)*lmax
		xs[n - 1] = n
        # We want the difference
		ys[n - 1] = fmlp_blocking - njlp_blocking

	print(xs)
	print(ys)

	cs = [c] * nmax
    # Specifically color this entry black for our figure to highlight it
	if (c == 'r'):
		cs[60] = '#400000'
	ax.bar(xs, ys, zs=k, zdir='y', color=cs, alpha=0.8)

ax.set_xlabel('n')
ax.set_ylabel('m')
ax.set_yticks(yticks, ms)
ax.set_zlabel('pi-blocking difference')

fig.savefig('plot.pdf', format='pdf')
plt.show()