import pdb
from numpy import *
NB = random.negative_binomial
from numpy.random import *

from pylab import *
check = lambda x, lo, hi: all(lo < y < hi for y in x)
def Mt(udot,n,k,iter):

    p0=0.5
    N0=udot+1
    #Prior on N = Poisson(mdot+1)
    #leads to conditional posterior on N-mdot =Poisson(lambda*prod(1-p_j))
    lam=udot
    N,p=N0,p0
    #prior on p =beta(1,1)
    #leads to  conditional posterior on p of beta(ndot+1,KN-ndot+1)
    N_trace=[]
    p_trace=[]
    it=[]
    U=N-udot
    #assume beta(alpha,beta)  a
    #all priors uniform here
    a=1.
    b=1.
    for i in range(iter):
        p=[beta(x+a,U+udot-x+b) for x in n]
        pi0=product([1-x for x in p])
        
        U=negative_binomial(udot+1,1.-pi0)
        #U=uniform(mdot,mdot/pi0)
        N_trace.append(U+udot)

       
        p_trace.append(p)
        it.append(i)
    figure(1)
    plot(it,N_trace)
    savefig('N')
    close()
   
    p_trace=transpose(p_trace)
    pstats={}
    for j in range(k):
        figure(j+2)
        plot(it,p_trace[j])     
        
        pstats[j]=stats_dept(p_trace[j][iter/2:iter-1])
        savefig('p'+str(j))
        close()


    
       
    Nstats=stats_dept(N_trace[iter/2:iter-1])
    return Nstats,pstats


def stats_dept(y):
    #return mean, se, and bci from trace
    x=sort(y)
    se=var(x)**0.5
    n=len(x)
    lower=int(0.025*n)
    upper=int(0.975*n)

    lci=x[lower]
    uci=x[upper]
    

    return mean(x),se,lci,uci




n=[27,23,26,22,23]
udot=54
t=5


N,p=Mt(udot,n,t,10000)


'N'
print 'mean', N[0]
print 'se', N[1]
print 'lci',N[2]
print 'uci',N[3]
'p'
for i in range(t):
    print 'p',i
    print 'mean', p[i][0]
    print 'se', p[i][1]
    print 'lci',p[i][2]
    print 'uci',p[i][3]

pdb.set_trace()