# Restuarant example for Decision Tree Learning, from Russell & Norvig 
# chapter 18

import sys

from example import Example
from attr_table import AttrTable
from dtl import DTL

def make_attr(alt, bar, fri, hun, pat, price, rain, res, type, est):
    return [alt, bar, fri, hun, pat, price, rain, res, type, est]

def make_example(alt, bar, fri, hun, pat, price, rain, res, type, est, wait):
    attr = make_attr(alt, bar, fri, hun, pat, price, rain, res, type, est)
    return Example(attr, wait)

table = AttrTable()

table.addBooleanAttribute('Alt', 0)
table.addBooleanAttribute('Bar', 1)
table.addBooleanAttribute('Fri', 2)
table.addBooleanAttribute('Hun', 3)
table.addAttribute('Pat', 4)
table.addAttribute('Price', 5)
table.addBooleanAttribute('Rain', 6)
table.addBooleanAttribute('Res', 7)
table.addAttribute('Type', 8)
table.addAttribute('Est', 9)
	
table.addAttrValue('Pat', 'None')
table.addAttrValue('Pat', 'Some')
table.addAttrValue('Pat', 'Full')
table.addAttrValue('Price', '$')
table.addAttrValue('Price', '$$')
table.addAttrValue('Price', '$$$')
table.addAttrValue('Type', 'French')
table.addAttrValue('Type', 'Thai')
table.addAttrValue('Type', 'Italian')
table.addAttrValue('Type', 'Burger')
table.addAttrValue('Est', '0-10')
table.addAttrValue('Est', '10-30')
table.addAttrValue('Est', '30-60')
table.addAttrValue('Est', '>60')

examples = []

examples.append(make_example('Yes', 'No', 'No', 'Yes', 'Some',  '$$$', 'No', 'Yes', 'French', '0-10', True))
examples.append(make_example('Yes', 'No', 'No', 'Yes', 'Full', '$',  'No', 'No', 'Thai', '30-60', False))
examples.append(make_example('No', 'Yes', 'No', 'No', 'Some', '$',  'No', 'No', 'Burger', '0-10', True))
examples.append(make_example('Yes', 'No', 'Yes', 'Yes', 'Full', '$', 'No', 'No', 'Thai', '10-30', True))
examples.append(make_example('Yes', 'No', 'Yes', 'No', 'Full',  '$$$', 'No', 'Yes','French',  '>60', False))
examples.append(make_example('No', 'Yes', 'No', 'Yes', 'Some', '$$', 'Yes', 'Yes', 'Italian', '0-10', True))
examples.append(make_example('No', 'Yes', 'No', 'No', 'None', '$',  'Yes', 'No', 'Burger', '0-10', False))
examples.append(make_example('No', 'No', 'No', 'Yes', 'Some', '$$',  'Yes', 'Yes', 'Thai', '0-10', True))
examples.append(make_example('No', 'Yes', 'Yes', 'No', 'Full', '$',  'Yes', 'No', 'Burger', '>60', False))
examples.append(make_example('Yes', 'Yes', 'Yes', 'Yes', 'Full',  '$$$', 'No', 'Yes','Italian','10-30',False))
examples.append(make_example('No', 'No', 'No', 'No', 'None', '$',  'No', 'No', 'Thai', '0-10', False))
examples.append(make_example('Yes', 'Yes', 'Yes', 'Yes', 'Full', '$', 'No', 'No', 'Burger', '30-60', True))

dtl = DTL(table)

#  I don't mind waiting
t = dtl.learn(examples, table.getAttributes(), True) 

# show the tree
print t

# test 'always wait if Patrons=Some' rule
print dtl.test(t, make_attr('', '', '',  '', 'Some', '', '', '', '', ''))

# test 'never wait if Patrons=Full and wait < 10 minutes' rule
print dtl.test(t, make_attr('', '', '',  '', 'Full', '', '', '', '', '0-10'))