import pdb
from PyMC import *
from numpy.random import *




def simdata(n,a,b,m):

    #simulate a vector of binomial outcomes with hetergeneous p
    
    success=[binomial(n,beta(a,b)) for i in range(m)]
    trials=[n]*m
    
    return trials,success,m


class m1(MetropolisHastings):

    def __init__(self,n,a,b,m):     
        'Initialization of model parameters and nodes' 
        MetropolisHastings.__init__(self)

        self.data=simdata(n,a,b,m)
        self.parameter('p',init_val=[0.5]*self.data[2],random=True)
        self.parameter('mu',init_val=0.0)
        self.parameter('tau',init_val=1.0)

       
    def model(self):      
        n,x,reps=self.data
       
        self.binomial_like(x,n,self.p)
        #likelihood on random p's
       
        self.normal_prior(logit(self.p),self.mu,self.tau)
        sd=1./sqrt(self.tau)
        self.half_normal_prior([sd],0.0001)
    
  
def run(iter=20000,burn=10000,thin=10,model='m1',plot=True,conv=True,auto=True,n=1000,reps=1,a=2.0,b=2.0,m=25):

    #the instance of the sampler is stored in the dictionary
    models={'m1':m1(n,a,b,m)}
    #now we invoke sampler reps times, each time simulating the data with specified vectors n, p
    for i in range(reps):
        #first time through model self initializes with simulated data
        sampler=models[model]
        sampler.sample(iterations=iter,burn=burn,thin=thin,plot=plot)

        results=sampler.summary(burn=burn,thin=thin)
        if conv:
            sampler.convergence(burn=burn,thin=thin)
        if auto:
            sampler.autocorrelation(burn=burn,thin=thin)
        name='stats'+str(i)
        sampler.export(results,name)

        #replaces data object with new simulated data for next iteration

        sampler.data=simdata(n,a,b,m)

if __name__=="__main__":
    run()