import sys
import seaborn as sns
sys.path.append('/home/fabee/code/')
from matplotlib.pyplot import *
from fabee.Plotting import *
from scipy import stats
from numpy import *

sns.set_context("talk", font_scale=1.5, rc={"lines.linewidth": 2.5})

def hinton(matrix, max_weight=None, ax=None):
    """Draw Hinton diagram for visualizing a weight matrix."""
    ax = ax if ax is not None else gca()

    if not max_weight:
        max_weight = 2**np.ceil(np.log(np.abs(matrix).max())/np.log(2))

    ax.patch.set_facecolor('gray')
    ax.set_aspect('equal', 'box')
    ax.xaxis.set_major_locator(NullLocator())
    ax.yaxis.set_major_locator(NullLocator())

    for (x,y),w in np.ndenumerate(matrix):
        color = 'white' if w > 0 else 'black'
        size = np.sqrt(np.abs(w))
        rect = Rectangle([x - size / 2, y - size / 2], size, size,
                             facecolor=color, edgecolor=color)
        ax.add_patch(rect)

    ax.autoscale_view()
    ax.invert_yaxis()

# ---------------------------------------------------------------------------
fig, ax = subplots()
fig.subplots_adjust(bottom=.2)

ax.bar([0,1],[.2,.8],facecolor='dodgerblue', alpha=.8,width=.7, align='center')
ax.set_title('Bernoulli distribution',fontsize=16, fontweight='bold')

ax.set_xlim([-.5,1.5])
box_off(ax)
#disjoint_axes(ax)
ax.set_xlabel('outcomes',fontsize=14, fontweight='bold')
ax.set_ylabel('P(outcome)',fontsize=14, fontweight='bold')
ax.set_xticks([0,1])
ax.set_xticklabels([0,1])
ax.set_ylim((0,1))

fig.savefig('figs/Bernoulli.pdf')

# ---------------------------------------------------------------------------
fig, ax = subplots()
fig.subplots_adjust(bottom=.2)
n = 5
k = arange(0,n)
ax.bar(k,0*k+1./n,facecolor='dodgerblue', alpha=.8,width=.7, align='center')
ax.set_title('uniform distribution',fontsize=16, fontweight='bold')

box_off(ax)
#disjoint_axes(ax)
ax.set_xlabel('k',fontsize=14, fontweight='bold')
ax.set_ylabel('P(X=k)',fontsize=14, fontweight='bold')
ax.set_xticks(k)
ax.set_xticklabels(k+1)
ax.set_ylim((0,1))
fig.savefig('figs/Uniform.pdf')

# ---------------------------------------------------------------------------

for i,(n,p) in enumerate(zip([10,20],[.5,.8])):
    fig, ax = subplots()

    fig.subplots_adjust(bottom=.2)
    k = arange(n+1)
    
    ax.bar(k,stats.binom.pmf(k,n,p),facecolor='dodgerblue', alpha=.8,width=.7, align='center')
    ax.set_title(r'binomial distribution $B\left(%.2f, %i\right)$' % (p,n),fontsize=16, fontweight='bold')
    
    box_off(ax)
    #disjoint_axes(ax)
    ax.set_xlabel('k',fontsize=14, fontweight='bold')
    ax.set_ylabel('P(k)',fontsize=14, fontweight='bold')
    ax.set_xticks(k)
    ax.set_xticklabels(k)
    ax.set_xlim((-1,n+1))
    ax.set_ylim((0,1))
    fig.savefig('figs/Binomial%02i.pdf' % (i,))


# ---------------------------------------------------------------------------
n = 20
for i, lam in enumerate([5, 0.05]):
    fig, ax = subplots()

    fig.subplots_adjust(bottom=.2)
    k = arange(n+1)
    
    ax.bar(k,stats.poisson.pmf(k,lam),facecolor='dodgerblue', alpha=.8,width=.7, align='center')
    ax.set_title(r'Poisson distribution $\lambda=%.2f$' % (lam,),fontsize=16, fontweight='bold')
    
    box_off(ax)
    #disjoint_axes(ax)
    ax.set_xlabel('k',fontsize=14, fontweight='bold')
    ax.set_ylabel('P(k)',fontsize=14, fontweight='bold')
    ax.set_xticks(k)
    ax.set_xticklabels(k)
    ax.set_xlim((-1,n+1))
    ax.set_ylim((0,1))
    fig.savefig('figs/Poisson%02i.pdf' % (i,))

# ---------------------------------------------------------------------------
fig, ax = subplots()

fig.subplots_adjust(bottom=.2)
t = linspace(-3,3,200)
ax.fill_between(t,stats.norm.pdf(t),facecolor='dodgerblue', alpha=.8)
ax.set_title(r'Gaussian/Normal distribution $N(\mu,\sigma)$',fontsize=16, fontweight='bold')
    
box_off(ax)
#disjoint_axes(ax)
ax.set_xlabel('x',fontsize=14, fontweight='bold')
ax.set_ylabel('p(x)',fontsize=14, fontweight='bold')
fig.savefig('figs/Gaussian00.pdf')

# ---------------------------------------------------------------------------

fig, ax = subplots()
n = 10
kk = 5
p = .5
fig.subplots_adjust(bottom=.2)
k = arange(n+1)
    
ax.bar(k,stats.binom.pmf(k,n,p),facecolor='dodgerblue', alpha=.8,width=.7, align='center')
ax.bar(k[:kk+1],stats.binom.pmf(k[:kk+1],n,p),facecolor='crimson', alpha=.5,width=.7, align='center')
ax.set_title(r'binomial distribution $B\left(\frac{1}{2}, %i\right)$' % (n,), fontsize=16, fontweight='bold')

box_off(ax)
#disjoint_axes(ax)
ax.set_xlabel('k',fontsize=14, fontweight='bold')
ax.set_ylabel('P(k)',fontsize=14, fontweight='bold')
ax.set_xticks(k)
ax.set_xticklabels(k)
ax.set_xlim((-1,n+1))
ax.set_ylim((0,1))
fig.savefig('figs/BinomialCdf00.pdf' )


fig, ax = subplots()
n = 10
kk = 5
p = .5
fig.subplots_adjust(bottom=.2)
k = arange(n+1)
    
ax.bar(k,stats.binom.pmf(k,n,p),facecolor='dodgerblue', alpha=.8,width=.7, align='center',label='p.m.f.')
ax.bar(k[:kk+1],stats.binom.pmf(k[:kk+1],n,p),facecolor='crimson', alpha=.5,width=.7, align='center')
ax.plot(k,stats.binom.cdf(k,n,p),'ok',mfc='crimson', alpha=1.,label='c.d.f.', ms=15)

ax.set_title(r'binomial distribution $B\left(\frac{1}{2}, %i\right)$' % (n,), fontsize=16, fontweight='bold')
ax.legend(frameon=False, loc='best')
box_off(ax)
#disjoint_axes(ax)
ax.set_xlabel('k',fontsize=14, fontweight='bold')
ax.set_ylabel('P(k)',fontsize=14, fontweight='bold')
ax.set_xticks(k)
ax.set_xticklabels(k)
ax.set_xlim((-1,n+1))
ax.set_ylim((0,1.1))
fig.savefig('figs/BinomialCdf01.pdf' )

fig, ax = subplots()
n = 10
kk = 2
p = .5
fig.subplots_adjust(bottom=.2)
k = arange(n+1)
    
ax.bar(k,stats.binom.pmf(k,n,p),facecolor='dodgerblue', alpha=.8,width=.7, align='center',label='p.m.f.')
ax.bar(k[:kk+1],stats.binom.pmf(k[:kk+1],n,p),facecolor='crimson', alpha=.5,width=.7, align='center')
ax.bar(k[-kk-1:],stats.binom.pmf(k[-kk-1:],n,p),facecolor='crimson', alpha=.5,width=.7, align='center')

ax.set_title(r'binomial distribution $B\left(\frac{1}{2}, %i\right)$' % (n,), fontsize=16, fontweight='bold')
ax.legend(frameon=False, loc='best')
box_off(ax)
#disjoint_axes(ax)
ax.set_xlabel('k',fontsize=14, fontweight='bold')
ax.set_ylabel('P(k)',fontsize=14, fontweight='bold')
ax.set_xticks(k)
ax.set_xticklabels(k)
ax.set_xlim((-1,n+1))
ax.set_ylim((0,1.1))
fig.savefig('figs/BinomialExample00.pdf' )


#------------------------------------------------------
fig = figure(figsize=(10,3.5))
ax = fig.add_axes([.1,.13,.6,.3])

n = 10
p = [.5,.8]
q = [.7, .3]
fig.subplots_adjust(bottom=0.2)
k = arange(n+1)
P = vstack((stats.binom.pmf(k,n,p[0])*q[0], stats.binom.pmf(k,n,p[1])*q[1])).T

hinton(P, ax = None)

#disjoint_axes(ax)
ax.set_xticks(k)
ax.set_xticklabels(k)
ax.set_yticks([0,1])
ax.set_ylim((-.5,1.5))
ax.set_xlim((-.5,n+.5))
ax.set_yticklabels(['subject #1', 'subject #2'])
fig.savefig('figs/Joint00.pdf' )

ax = fig.add_axes([.75,.13,.2,.3])
ax.barh([0,1],q, facecolor='dodgerblue',alpha=.8, align='center')
box_off(ax)
#disjoint_axes(ax)
ax.set_xticks([0,.5,1.])
ax.set_yticks([])
ax.set_ylim((-.5,1.5))
fig.savefig('figs/Joint01.pdf' )

ax = fig.add_axes([.1,.6,.6,.2])
ax.bar(k,sum(P,axis=1), facecolor='dodgerblue',alpha=.8, align='center')
a = .7
ax.axis([-a,n-a+1.5,0,1])
box_off(ax)
#disjoint_axes(ax)
ax.set_xticks([])
ax.set_yticks([0,.3])
ax.set_ylim((0,.3))
fig.savefig('figs/Joint02.pdf' )


#------------------------------------------------------
n = 10
k = arange(n+1)
p = [.5,.8]
q = [.7, .3]
P = vstack((stats.binom.pmf(k,n,p[0])*q[0], stats.binom.pmf(k,n,p[1])*q[1]))
Pk = sum(P,axis=0)

fig = figure()
for i,kk in enumerate(k):
    ax = fig.add_subplot(3,4,i+1)
    fig.subplots_adjust(bottom=0.2)

    ax.bar([0,1],P[:,i]/Pk[i], facecolor='dodgerblue',alpha=.8, align='center')

    
    #disjoint_axes(ax)
    ax.set_xticks([0,1])
    ax.set_xticklabels(['#1','#2'], fontsize=8)
    ax.set_yticks([0,.5,1])
    ax.set_yticklabels([0,.5,1],fontsize=8)
    ax.set_xlim((-.5,1.5))
    ax.set_ylim((0,1))
    ax.set_title('P({#1,#2}| %i successes)'  % (i,), fontsize=8)

fig.subplots_adjust(wspace=.8, hspace=.8)
fig.savefig('figs/Posterior00.pdf')