# -*- coding: utf-8 -*-

"""

Created on Mon May 21 14:39:15 2018

@author: lijie

"""

import matplotlib.pyplot as plt

'''plot the tree'''

#Set the type of text box and the arrow(箭头)

decisionNode=dict(boxstyle='sawtooth',fc='0.8')

leafNode = dict(boxstyle="round4", fc="0.8")

arrow_args = dict(arrowstyle="<-")

#plot the arrow with comment

def plotNode(nodeTxt, centerPt, parentPt, nodeType):

createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',

xytext=centerPt, textcoords='axes fraction',

va="center", ha="center", bbox=nodeType, arrowprops=arrow_args )

#plot a node

def createPlot(inTree):

fig = plt.figure(1, facecolor='white')

fig.clf()

axprops= dict(xticks=[],yticks=[])

createPlot.ax1 = plt.subplot(111, frameon=False,**axprops)

plotTree.totalW=float(getNumLeafs(inTree))

plotTree.totalD=float(getTreeDepth(inTree))

plotTree.xoff=-0.5/plotTree.totalW

plotTree.yoff=1.0

plotTree(inTree,(0.5,1.0),'')

#plotNode('decision node', (0.5, 0.1), (0.1, 0.5), decisionNode)

#plotNode('leaf node', (0.8, 0.1), (0.3, 0.8), leafNode)

plt.show()

#judge if the nodes are dictonaires, if not → they are leaf nodes

def getNumLeafs(myTree):

numLeafs = 0

firstSide = list(myTree.keys())

firstStr = firstSide[0]

secondDict = myTree[firstStr]

for key in secondDict.keys():

if type(secondDict[key]).__name__=='dict':

numLeafs += getNumLeafs(secondDict[key])

else: numLeafs +=1

return numLeafs

#judge if the nodes are dictonaires, if not → they are leaf nodes

def getTreeDepth(myTree):

maxDepth = 0

firstSide = list(myTree.keys())

firstStr = firstSide[0]

secondDict = myTree[firstStr]

for key in secondDict.keys():

if type(secondDict[key]).__name__=='dict':

thisDepth = 1 + getTreeDepth(secondDict[key])

else: thisDepth = 1

if thisDepth > maxDepth: maxDepth = thisDepth

return maxDepth

#Store the tree information in advance.save time

def retrieveTree(i):

listOfTrees =[{'no surfacing': {0: 'no', 1: {'flippers': {0: 'no', 1: 'yes'}}}},

{'no surfacing': {0: 'no', 1: {'flippers': {0: {'head': {0: 'no', 1: 'yes'}}, 1: 'no'}}}}

]

return listOfTrees[i]

def plotMidText(cntrPt, parentPt, txtString):

xMid = (parentPt[0]-cntrPt[0])/2.0 + cntrPt[0]

yMid = (parentPt[1]-cntrPt[1])/2.0 + cntrPt[1]

createPlot.ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)

#if the first key tells you what feat was split on

def plotTree(myTree, parentPt, nodeTxt):

numLeafs = getNumLeafs(myTree) #x: width of this tree

depth = getTreeDepth(myTree)

firstSide = list(myTree.keys())

firstStr = firstSide[0] #the text label for this node should be this

cntrPt = (plotTree.xoff + (1.0 + float(numLeafs))/2.0/plotTree.totalW, plotTree.yoff)

plotMidText(cntrPt, parentPt, nodeTxt)

plotNode(firstStr, cntrPt, parentPt, decisionNode)

secondDict = myTree[firstStr]

plotTree.yoff = plotTree.yoff - 1.0/plotTree.totalD

for key in secondDict.keys():

#test if the nodes are dictonaires, if not → leaf nodes

if type(secondDict[key]).__name__=='dict':

plotTree(secondDict[key],cntrPt,str(key)) #recursion

#if it is a leaf node → print

else:

plotTree.xoff = plotTree.xoff + 1.0/plotTree.totalW

plotNode(secondDict[key], (plotTree.xoff, plotTree.yoff), cntrPt, leafNode)

plotMidText((plotTree.xoff, plotTree.yoff), cntrPt, str(key))

plotTree.yoff = plotTree.yoff + 1.0/plotTree.totalD

if __name__ == '__main__':

myTree=retrieveTree(0)

createPlot(myTree)