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

# ## Examples of using the Bayesian-belief-networks library by ebay
# 
# #### https://github.com/eBay/bayesian-belief-networks

# In[2]:


# Burglar and racoon

from bayesian.factor_graph import *


def fP(P):
    '''Police'''
    if P is True:
        return 1
    elif P is False:
        return 0


def fT(T):
    '''Trash Can fell'''
    if T is True:
        return 0
    elif T is False:
        return 1


def fDBR(D, B, R):
    '''Dog bark, burglar, racoon'''
    table = dict()
    table['ttt'] = 0.002
    table['ttf'] = 0.008
    table['tft'] = 0.099
    table['tff'] = 0.0792
    table['ftt'] = 0
    table['ftf'] = 0
    table['fft'] = 0.099
    table['fff'] = 0.7128
    key = ''
    key = key + 't' if D else key + 'f'
    key = key + 't' if B else key + 'f'
    key = key + 't' if R else key + 'f'
    return table[key]


def fPD(P, D):
    '''Police came given dog bark'''
    table = dict()
    table['tt'] = 0.4
    table['tf'] = 0.01
    table['ft'] = 0.6
    table['ff'] = 0.99
    key = ''
    key = key + 't' if P else key + 'f'
    key = key + 't' if D else key + 'f'
    return table[key]


def fTR(T, R):
    '''Trash can fall given Racoon came'''
    table = dict()
    table['tt'] = 0.7
    table['tf'] = 0.05
    table['ft'] = 0.3
    table['ff'] = 0.95
    key = ''
    key = key + 't' if T else key + 'f'
    key = key + 't' if R else key + 'f'
    return table[key]

if __name__ == '__main__':
    g = build_graph(
        fDBR,fPD,fTR)
    
#     g.q()
#     g.q(P='high')
#     g.q(D=True)
#     g.q(S=True)
#     g.q(C=True, S=True)
#     g.q(D=True,R=False)
    g.q(P=True,T=True)


# In[3]:


# cancer example

'''This is the example from Chapter 2 BAI'''
from bayesian.factor_graph import *


def fP(P):
    '''Pollution'''
    if P == 'high':
        return 0.1
    elif P == 'low':
        return 0.9


def fS(S):
    '''Smoker'''
    if S is True:
        return 0.3
    elif S is False:
        return 0.7


def fC(P, S, C):
    '''Cancer'''
    table = dict()
    table['ttt'] = 0.05
    table['ttf'] = 0.95
    table['tft'] = 0.02
    table['tff'] = 0.98
    table['ftt'] = 0.03
    table['ftf'] = 0.97
    table['fft'] = 0.001
    table['fff'] = 0.999
    key = ''
    key = key + 't' if P == 'high' else key + 'f'
    key = key + 't' if S else key + 'f'
    key = key + 't' if C else key + 'f'
    return table[key]


def fX(C, X):
    '''X-ray'''
    table = dict()
    table['tt'] = 0.9
    table['tf'] = 0.1
    table['ft'] = 0.2
    table['ff'] = 0.8
    key = ''
    key = key + 't' if C else key + 'f'
    key = key + 't' if X else key + 'f'
    return table[key]


def fD(C, D):
    '''Dyspnoeia'''
    table = dict()
    table['tt'] = 0.65
    table['tf'] = 0.35
    table['ft'] = 0.3
    table['ff'] = 0.7
    key = ''
    key = key + 't' if C else key + 'f'
    key = key + 't' if D else key + 'f'
    return table[key]


if __name__ == '__main__':
    g = build_graph(
        fP, fS, fC, fX, fD,
        domains={
            'P': ['low', 'high']})
#     g.q()
#     g.q(P='high')
#     g.q(D=True)
#     g.q(S=True)
#     g.q(C=True, S=True)
    g.q(D=True, S=True)


# In[ ]: