import numpy as np
from sklearn import tree

questions = ('《Python程序设计与数据采集》',
             '《Python数据分析、挖掘与可视化》',
             '《Python程序设计基础（第3版）》',
             '《Python程序设计基础与应用（第2版）》',
             '《Python程序设计实例教程（第2版）》',
             '《Python程序设计（第3版）》', '《Python网络程序设计》',
             '《Python数据分析与数据可视化》',
             '《大数据的Python基础（第2版）》',
             '《Python程序设计开发宝典》',
             '《Python可以这样学》', '《中学生可以这样学Python》',
             '《Python编程基础与案例集锦（中学版）》',
             '《Python程序设计实验指导书》',
             '《Python程序设计入门与实践》',
             '《Python程序设计实用教程》',
             '微信公众号“Python小屋”的免费资料')
# 每个样本的数据含义：0没看过，1很多看不懂，2大部分可以看懂，3没压力
answers = [[3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
           [1, 0, 2, 0, 0, 0, 0, 0, 3, 1, 0, 0, 0, 2, 2, 2, 1],
           [2, 2, 3, 3, 3, 1, 1, 1, 3, 2, 2, 3, 3, 3, 0, 3, 2],
          
           [2, 2, 3, 2, 2, 1, 1, 3, 1, 2, 2, 3, 3, 2, 2, 2, 1],
           [3, 3, 3, 0, 3, 2, 2, 3, 3, 2, 3, 3, 3, 3, 3, 3, 2],
           [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3],
           [0, 2, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
           [2, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]
          ]
labels = ['超级高手', '门外汉', '初级选手', '初级选手', '高级选手',
          '中级选手', '高级选手', '超级高手', '初级选手', '初级选手']

clf = tree.DecisionTreeClassifier().fit(answers, labels) # 训练

yourAnswer = []
print('在下面的问卷中，'
      '没看过输入0，很多看不懂输入1，大部分可以看懂输入2，没压力输入3')
# 显示调查问卷，并接收用户输入
for question in questions:
    print('=========\n你看过董付国老师的', question, '吗？')
    # 确保输入有效
    while True:
        try:
            answer = int(input('请输入：'))
            assert 0<=answer<=3
            break
        except:
            print('输入无效，请重新输入。')
            pass
    yourAnswer.append(answer)
    
print(clf.predict(np.array(yourAnswer).reshape(1,-1)))   # 分类