#!/usr/bin/env python
# coding: utf-8

# In[7]:


from pymc3 import Categorical, Model, DensityDist
from pymc3.distributions.dist_math import bound
from theano.tensor import switch, log, le, constant
from theano import shared, gradient
import numpy as np

with Model() as simple_model:
    
    Age = Categorical('Age', p=[0.3, 0.5], testval=1)
    
    Sex = Categorical('Sex', p=[0.6, 0.4], testval=0)
    
    def edu_logp(value):
        p = constant(np.array([[0.25, 0.28, 0.12], [0.36, 0.30, 0.10]]))
        return bound(switch(value, log(p[Sex, Age]), log(1 - p[Sex, Age])),
            value >= 0, value <= 1)

    Education = DensityDist('Education', edu_logp, dtype='int64', testval=0)
    
    def res_logp(value):
        p = constant(np.array([0.75, 0.8]))
        return bound(switch(value, log(p[Education]), log(1 - p[Education])),
            value >= 0, value <= 1)

    Residence = DensityDist('Residence', res_logp, dtype='int64', testval=0)
    
    
    def occ_logp(value):
        p = constant(np.array([0.04, 0.08]))
        return bound(switch(value, log(p[Education]), log(1 - p[Education])),
            value >= 0, value <= 1)

    Occupation = DensityDist('Occupation', occ_logp, dtype='int64', testval=0)
    
    
    def travel_logp(value):
        p =  constant(np.array([[[0.48, 0.42, 0.10],  # Small and employer
                              [0.56, 0.36, 0.08]], # Small and self-employed
                             [[0.58, 0.24, 0.18],  # Big and employer
                              [0.70, 0.21, 0.09]]])) # Big and self-employed

        return bound(log(p[Residence, Occupation, value]),
            value >= 0,
            value <= 2)

    Travel = DensityDist('Travel', travel_logp, dtype='int64', testval=0)


# In[8]:


from pymc3 import find_MAP, sample

with simple_model:
    tr = sample(1000)


# In[ ]: