import math,random
import matplotlib.pyplot as plt

# Binary Entropy function
def BinEnt(p):
  if p <= 0.0 or p >= 1.0:
    return 0.0
  return p*math.log2(1.0/p) + (1.0-p)*math.log2(1.0/(1.0-p))


# https://wiki.python.org/moin/BitManipulation
def ones_and_length(n):
  ones = 0
  length = 0
  while n > 0:
    ones += (n & 1)
    length += 1
    n >>= 1
  return ones, length

# return entropy with leading-style length and non-leading
def bit_ents(n, full_length):
  ones, length = ones_and_length(n)
  if ones == 0:
    return 0.0, 0.0
  leading_p = float(ones)/float(length)
  non_p = float(ones)/float(full_length)
  return BinEnt(leading_p), BinEnt(non_p)


def trial(full_length):
  maximum = (1 << full_length) - 1
  n = random.randint(0, maximum)
  m = random.randint(0, maximum)
  lHn,Hn = bit_ents(n, full_length)
  lHtimes, Htimes = bit_ents(n*m, 2*full_length)
  lHplus, Hplus = bit_ents(n+m, full_length+1)
  return [lHn, lHtimes, lHplus, Hn, Htimes, Hplus]


def do_trials(full_length, num_trials):
  results = [0.0]*6
  for t in range(num_trials):
    temp = trial(full_length)
    for j in range(len(temp)):
      results[j] += temp[j]
  for j in range(len(results)):
    results[j] /= float(num_trials)
  return results


num_trials = 10000
lengths = [10, 20, 40, 80, 160, 320, 640, 1280]
#lengths = [5, 10, 20, 40, 80]
res = [do_trials(l, num_trials) for l in lengths]
addone_res = [do_trials(l+1, num_trials) for l in lengths]

plt.figure()
plt.plot(lengths, [a[0] for a in res], label="$G(k)$")
plt.plot(lengths, [a[2] for a in res], label="$H_{n+m}$")
plt.plot(lengths, [a[1] for a in res], label="$H_{n\cdot m}$")
plt.title("Bit 'entropy'")
plt.xlabel("k")
plt.ylabel("average entropy")
plt.legend()


plt.figure()
plt.plot(lengths, [a[0] for a in addone_res], label="$G(k+1)$")
plt.plot(lengths, [a[2] for a in res], label="$H_{n+m}$")
plt.title("Bit 'entropy'")
plt.xlabel("k")
plt.ylabel("average entropy")
plt.legend()

plt.figure()
plt.plot(lengths[:-1], [a[0] for a in res][1:], label="$G(2k)$")
plt.plot(lengths[:-1], [a[1] for a in res][:-1], label="$H_{n\cdot m}$")
plt.title("Bit 'entropy'")
plt.xlabel("k")
plt.ylabel("average entropy")
plt.legend()


plt.show()