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

from matplotlib.ticker import NullFormatter

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

# --------------- PLOT 1 -------------------------
# the random data
distr = stats.uniform
col = '+*0<>v'

for k,distr in enumerate([stats.laplace, stats.norm, stats.expon,stats.uniform]):
    col = [col[i] for i in random.permutation(6)]
    x = random.randn(5000)

    nullfmt   = NullFormatter()         # no labels

    # definitions for the axes
    left, width = 0.1, 0.65
    bottom, height = 0.1, 0.65
    bottom_h = left_h = left+width+0.02
    
    rect_scatter = [left + 0.22, bottom + 0.22 , width, height]
    rect_histx = [left + 0.22, bottom, width, 0.2]
    rect_histy = [left, bottom + 0.22 , 0.2, height]
    
    # start with a rectangular Figure
    fig = figure(figsize=(8,8))
    
    axQQ = axes(rect_scatter)
    axHistx = axes(rect_histx)
    axHisty = axes(rect_histy)
    
    # no labels
    axHistx.yaxis.set_major_formatter(nullfmt)
    axHisty.xaxis.set_major_formatter(nullfmt)
    axQQ.xaxis.set_major_formatter(nullfmt)
    axQQ.yaxis.set_major_formatter(nullfmt)
    
    
    
    # the scatter plot:
    z = distr.ppf(stats.norm.cdf(x))
    y = linspace(amin(z),amax(z),1000)
    
    
    z = distr.ppf(stats.norm.cdf(x))
    if distr != stats.norm:
        if distr == stats.uniform:
            axQQ.plot(x, z,'ok',marker=col[0],ms=5,label='c.d.f.')
        else:
            axQQ.plot(x, z,'ok',marker=col[0],ms=5,label='correct')

        if distr != stats.expon:
            axQQ.plot((z-amin(z))/(amax(z)-amin(z))*(amax(x)-amin(x)) + amin(x),\
                      (x-amin(x))/(amax(x)-amin(x))*(amax(z)-amin(z)) + amin(z),'ok',marker=col[1],ms=5)
            axQQ.plot(x, (x-amin(x))/(amax(x)-amin(x))*(amax(z)-amin(z)) + amin(z),'ok',marker=col[2],ms=5)
        
    
    # now determine nice limits by hand:
    axHistx.hist(x, bins=100,normed=True)
    if distr != stats.expon:
        axHisty.plot(distr.pdf(y),y)
        z2 = distr.pdf(y)
        y = hstack((y[0],y,y[-1]))
        z2 = hstack((0,z2,0))
        axHisty.fill(z2,y,color=(.0,.0,1.))
    
    axQQ.set_xlim(axHistx.get_xlim())
    axQQ.set_ylim(axHisty.get_ylim())

    if distr == stats.uniform:
        axQQ.set_ylim((-.1,1.1))
        axHisty.set_ylim((-.1,1.1))
        axHisty.set_xlim((.0,1.1))
    
    axHistx.set_xlabel('x',fontsize=16)
    axHistx.set_ylabel('p(x)',fontsize=16)
    axHisty.set_ylabel('y',fontsize=16)
    axHisty.set_xlabel('p(y)',fontsize=16)

    fig.savefig('figs/HE%i.png' % (k,))
    if distr == stats.norm:
        axQQ.plot(x, z,'ok',marker=col[0],ms=5)
    elif distr == stats.expon:
        axHisty.plot(distr.pdf(y),y)
        z2 = distr.pdf(y)
        y = hstack((y[0],y,y[-1]))
        z2 = hstack((0,z2,0))
        axHisty.fill(z2,y,color=(.0,.0,1.))
       
    else:
        axQQ.legend(loc=2)
    fig.savefig('figs/HE%iSolution.png' % (k,))

# ####################################################3
fig = figure()


ax = fig.add_subplot(111)
xx = linspace(-3.,stats.norm.ppf(1-0.2),1000)

x = linspace(-3.,3.,1000)
y = stats.norm.pdf(x,scale=1)
yy = stats.norm.pdf(xx,scale=1)
yy[0] = 0
yy[-1] = 0

ax.plot(x,y,'k-',lw=2)
ax.plot(x,stats.norm.pdf(x),'k-',lw=1)
ax.set_xlabel('x',fontsize=16)
ax.set_ylabel('pdf',fontsize=16)
ax.fill(xx,yy,'b')

ax.set_xlim(-3.,3.)

ax.text(xx[-1],-.1,'b');

ax.text(xx[-1],.4,'p(x)',color='k');
ax.text(xx[0],.3,'F(b) = P(x <= b)',color='b');

#XKCDify(ax, expand_axes=True,yaxis_loc=0,xaxis_loc=0)

fig.savefig('figs/cdf.png')

#-----------------------------
fig = figure()


ax = fig.add_subplot(111)
xx = linspace(-3.,stats.norm.ppf(1-0.2),1000)

x = linspace(-3.,3.,1000)
y = stats.norm.pdf(x,scale=1)
yy = stats.norm.pdf(xx,scale=1)
yy[0] = 0
yy[-1] = 0

ax.plot(x,y,'k-',lw=2)
ax.plot(x,stats.norm.cdf(x),'b-',lw=1)
ax.set_xlabel('x/b',fontsize=16)
ax.set_ylabel('pdf/cdf',fontsize=16)

ax.set_xlim(-3.,3.)

ax.text(xx[-1],.4,'p(x)',color='k');
ax.text(xx[0],.3,'F(b) = P(x <= b)',color='b');

#XKCDify(ax, expand_axes=True,yaxis_loc=0,xaxis_loc=0)

fig.savefig('figs/cdf2.png')

# ####################################################3
fig = figure()


ax = fig.add_subplot(111)

x = hstack((linspace(-3.,stats.norm.ppf(0.13),1000),\
            linspace(stats.norm.ppf(1-0.13),3.,1000)))

xx = hstack((linspace(-3.,stats.norm.ppf(0.2),1000),\
            linspace(stats.norm.ppf(1-0.2),3.,1000)))

y = stats.norm.pdf(x,scale=1)
yy = stats.norm.pdf(xx,scale=1)

y[[0,999,1000,-1]] = 0
yy[[0,999,1000,-1]] = 0

t = linspace(-3.,3.,1000)
ax.plot(t,stats.norm.pdf(t),'k-',lw=2)

ax.fill(xx[:1000],yy[:1000],'b')
ax.fill(xx[1000:],yy[1000:],'b')
ax.text(xx[1000],-.1,'b')
ax.text(xx[999],-.1,'-b')
ax.text(.2,.7,'P(|x|>b) =$\\alpha$',color='b');

#XKCDify(ax, expand_axes=True,yaxis_loc=0,xaxis_loc=0)

fig.savefig('figs/pval0.png')

#---------------------------------------------------
fig = figure()

ax = fig.add_subplot(111)

t = linspace(-3.,3.,1000)
ax.plot(t,stats.norm.pdf(t),'k-',lw=2)


ax.fill(x[:1000],y[:1000],'r')
ax.fill(x[1000:],y[1000:],'r')


ax.text(x[1000],-.1,'t')
ax.text(x[999],-.1,'-t')

ax.text(.2,.5,'P(|x| > t) = p-value',color='r');

#XKCDify(ax, expand_axes=True,yaxis_loc=0,xaxis_loc=0)

fig.savefig('figs/ pval1.png')



# show()
                      

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