#!/usr/bin/python

"""
Title: Entropy simulator
Purpose: To estimate the mutual information of a random channel model by 
simulating and collecting the channels statistics over sufficient iterations.
Author: Chris Nicola
Version: 0.1 (alpha)
"""

#Import
from scipy import *
from pylab import *
from pickle import *
from scipy.stats import norm
from scipy.stats import bernoulli

#Functions

def capacity(N, pwr, Ps, corr):
	PsT_sR = corr + Ps*(1-corr)
	sum = 0.
	sd_x = sqrt(pwr/Ps)
	sd_xn = sqrt(pwr/Ps + 1)
	for i in range(N):
		s_T = bernoulli.rvs(PsT_sR)
		n = norm.rvs(size=2)
		x = norm.rvs(size=2, loc=0.0, scale=sd_x)
		y = s_T*x + n
		mut_info = log(prob_Y_X_Sr1(y, x, PsT_sR)/prob_Y_Sr1(y, sd_xn, PsT_sR))
		sum = sum + mut_info/N
	return sum*Ps/log(2)

def prob_Y_Sr1(y, sd_xn, PsT_sR):
	a = norm.pdf(y[0], loc=0., scale=sd_xn)*norm.pdf(y[1], loc=0., scale=sd_xn)
	b = norm.pdf(y[0])*norm.pdf(y[1])
	return PsT_sR*a + (1-PsT_sR)*b
	
def prob_Y_X_Sr1(y, x, PsT_sR):
	a = norm.pdf(y[0]-x[0])*norm.pdf(y[1]-x[1])
	b = norm.pdf(y[0])*norm.pdf(y[1])
	return PsT_sR*a + (1-PsT_sR)*b

print "Computing Capacity:"
c = range(11)
for i in range(11):
  c[i] = capacity(1000000, 10, 0.9, i/10.)
  print c[i]

x = [0., .1, .2, .3, .4, .5, .6, .7, .8, .9, 1.]
plot(x, c, 'bx-')
xlabel(r'$Correlation of S_R and S_T$')
ylabel('Mutual information in bps/Hz')
title(r'$Model 1: Y=S_R(S_T \dot X + N)) with P_S=0.9$')
savefig("capacity_0.9_10dB.png")
# Pickle the results for later use
file = open('results.pycap', 'w')
pickle.dump(c, file)
pickle.dump(x, file)
print "Simulation is finished, results saved"	
show()