import pdb
#import profile
# Import relevant modules
from numpy import  *
from PyMC import  *


def simdata(N0,years,s,py,cv,surv_n,age_n):

    N=[N0]


    surv=[]
    young=[]
    
   

    for y in range(years):

        

        f=py/(1-py)

        N.append(binomial(N[y],s)+poisson(N[y]*f))

        surv.append(binomial(surv_n,s))
        young.append(binomial(age_n,py))

        
        

    Nhat=[]
    Nhat_se=[]
    

    for y in range(years+1):
        se=N[y]*cv
        Nhat.append(normal(N[y],se))
        Nhat_se.append(se)
        
        
        

    print 'true states', N
    print 'surveys', Nhat,Nhat_se
    print 'age data', age_n,young
    print 'survival data',surv_n,surv


    return years,N,Nhat,Nhat_se,young,age_n,surv_n,surv

def data():
    surveys=[1075, 960, 1171, 981, 1049, 1091, 1056, 1160, 1150, 1370, 1210]
    nyears=10
    young=[99, 101, 100, 102, 124, 98, 124, 84, 110, 102]
    age_sample=100
    surv_sample=100
    surv=[53, 54, 55, 40, 50, 49, 45, 52, 41, 45]

    return(nyears,surveys,young,age_sample,surv_sample,surv)




class estim(MetropolisHastings):

    def __init__(self,N0,years,s,py,cv,surv_n,age_n):        
        'Initialization of model parameters and nodes' 
        MetropolisHastings.__init__(self)
        self.data=simdata(N0,years,s,py,cv,surv_n,age_n)
        self.years=self.data[0]
        self.N=self.data[1]
        self.Nhat=self.data[2]

        self.Nhat_se=self.data[3]

        self.young=self.data[4]
        self.age_n=self.data[5]
        self.surv_n=self.data[6]
        self.survive=self.data[7]
          


        init_py=[float(y)/self.age_n for y in self.young]

        init_s=[float(y)/self.surv_n for y in self.survive]


        self.parameter('py',init_val=init_py)
        
        self.parameter('s',init_val=init_s)

        
        self.parameter('N0',init_val=1000.)
        self.node('N',(self.years+1,))


             
        
    def model(self):       
        
        self.uniform_prior(self.N0,500,2000)

        self.uniform_prior(self.py,0,1)

        self.uniform_prior(self.s,0,1)

      
        
##
        self.N[0]=self.N0
##
##        
##        
        self.binomial_like(self.young,[self.age_n]*self.years,self.py)


        self.binomial_like(self.survive,[self.surv_n]*self.years,self.s)

     


        

        for y in range(self.years):


            try:
                survivors=binomial(self.N[y],self.s[y])
                f=self.py[y]/(1.-self.py[y])
                recruits=poisson(self.N[y]*f)
  
                self.N[y+1]=survivors+recruits
            except ValueError:
                pass
##
##                  
##            
##
        tau=[1./y**2 for y in self.Nhat_se]
##
##        
        self.normal_like(self.Nhat,self.N,tau)
##
##
   

   

class estim2(MetropolisHastings):

    def __init__(self,N0,years,s,py,cv,surv_n,age_n):        
        'Initialization of model parameters and nodes' 
        MetropolisHastings.__init__(self)
        self.data=simdata(N0,years,s,py,cv,surv_n,age_n)
       
        self.years=self.data[0]
        self.N=self.data[1]
        self.Nhat=self.data[2]

        self.Nhat_se=self.data[3]

        self.young=self.data[4]
        self.age_n=self.data[5]
        self.surv_n=self.data[6]
        self.survive=self.data[7]
          


        init_py=mean([float(y)/self.age_n for y in self.young])

        init_s=mean([float(y)/self.surv_n for y in self.survive])


        self.parameter('py',init_val=init_py)
        
        self.parameter('s',init_val=init_s)

        
        self.parameter('N0',init_val=1000.)
        self.node('N',(self.years+1,))


             
        
    def model(self):       
        
        self.uniform_prior(self.N0,500,2000)

        self.uniform_prior(self.py,0,1)

        self.uniform_prior(self.s,0,1)

      
        
##
        self.N[0]=self.N0
##
##        
##        
        self.binomial_like(self.young,[self.age_n]*self.years,[self.py]*self.years)


        self.binomial_like(self.survive,[self.surv_n]*self.years,[self.s]*self.years)

     


        

        for y in range(self.years):


            try:
                survivors=binomial(self.N[y],self.s)


                f=self.py/(1.-self.py)
                recruits=poisson(self.N[y]*f)

  
                self.N[y+1]=survivors+recruits
            except ValueError:
                pass
##
##                  
##            
##
        tau=[1./y**2 for y in self.Nhat_se]
##
##        
        self.normal_like(self.Nhat,self.N,tau)
##
##
   

   


        


            

            
   

      


def run(iter=10000,burn=5000,thin=1,plot=True):
  # print simulate_data(N0,CV,P_m,phi,f,years,f_sample,s_sample)
    sampler=estim2(1000,10,.5,.3333,0.10,10,10)
    results = sampler.sample(iterations = iter, burn = burn, thin = thin, plot = plot)

    sampler.export(results,'stats')

if __name__  ==  '__main__':
   run()