import sympy
from sympy import stats, simplify, Rational, Integer,Eq
from sympy.stats import density, E
from sympy.abc import a

x=stats.Die('D1',6) # 1st six sided die
y=stats.Die('D2',6) # 2nd six sides die
z = x+y             # sum of 1st and 2nd die

J = stats.E((x - a*(x+y))**2)         # expectation
sol=sympy.solve(sympy.diff(J,a),a)[0] # using calculus to minimize
print sol # solution is 1/2

fig, ax = subplots()
v = arange(1,7) + arange(1,7)[:,None]
foo=lambda i: density(z)[Integer(i)].evalf() # some tweaks to get a float out
Zmass=array(map(foo,v.flat),dtype=float32).reshape(6,6)

ax.pcolor(arange(1,8),arange(1,8),Zmass,cmap=cm.gray)
ax.set_xticks([(i+0.5) for i in range(1,7)])
ax.set_xticklabels([str(i) for i in range(1,7)])
ax.set_yticks([(i+0.5) for i in range(1,7)])
ax.set_yticklabels([str(i) for i in range(1,7)])
for i in range(1,7):
    for j in range(1,7):
        ax.text(i+.5,j+.5,str(i+j),fontsize=18,color='y')
ax.set_title(r'Probability Mass for $Z$',fontsize=18)    
ax.set_xlabel('$X$ values',fontsize=18)
ax.set_ylabel('$Y$ values',fontsize=18);

#generate samples conditioned on z=7
samples_z7 = lambda : stats.sample(x, sympy.Eq(z,7)) # Eq constrains Z
mn= mean([(6-samples_z7())**2 for i in range(100)]) #using 6 as an estimate
mn0= mean([(7/2.-samples_z7())**2 for i in range(100)]) #7/2 is the MSE estimate
print 'MSE=%3.2f using 6 vs MSE=%3.2f using 7/2 ' % (mn,mn0)

# here 6 is ten times more probable than any other outcome
x=stats.FiniteRV('D3',{1:Rational(1,15), 2:Rational(1,15), 3: Rational(1,15), 
                       4:Rational(1,15), 5:Rational(1,15), 6: Rational(2,3)})
z = x + y

# now re-create the plot
fig, ax = subplots()
foo=lambda i: density(z)[Integer(i)].evalf() # some tweaks to get a float out
Zmass=array(map(foo,v.flat),dtype=float32).reshape(6,6)

ax.pcolor(arange(1,8),arange(1,8),Zmass,cmap=cm.gray)
ax.set_xticks([(i+0.5) for i in range(1,7)])
ax.set_xticklabels([str(i) for i in range(1,7)])
ax.set_yticks([(i+0.5) for i in range(1,7)])
ax.set_yticklabels([str(i) for i in range(1,7)])
for i in range(1,7):
    for j in range(1,7):
        ax.text(i+.5,j+.5,str(i+j),fontsize=18,color='y')
ax.set_title(r'Probability Mass for $Z$; Nonuniform case',fontsize=16)    
ax.set_xlabel('$X$ values',fontsize=18)
ax.set_ylabel('$Y$ values',fontsize=18);

E(x, Eq(z,7)) # conditional expectation E(x|z=7)

#generate samples conditioned on z=7
samples_z7 = lambda : stats.sample(x, Eq(z,7)) # Eq constrains Z
mn= mean([(6-samples_z7())**2 for i in range(100)]) #using 6 as an estimate
mn0= mean([(5-samples_z7())**2 for i in range(100)]) #7/2 is the MSE estimate
print 'MSE=%3.2f using 6 vs MSE=%3.2f using 5 ' % (mn,mn0)