from collections import Counter
import pandas as pd # 数据处理
import numpy as np # 数学运算
from sklearn.model_selection import train_test_split, cross_validate, GridSearchCV # 划分数据集函数
from sklearn.metrics import accuracy_score # 准确率函数
from sklearn.tree import DecisionTreeClassifier # sklearn的决策树模型
RANDOM_SEED = 2020 # 固定随机种子

csv_data = './data/high_diamond_ranked_10min.csv' # 数据路径
data_df = pd.read_csv(csv_data, sep=',') # 读入csv文件为pandas的DataFrame
data_df = data_df.drop(columns='gameId') # 舍去对局标号列

print(data_df.iloc[0]) # 输出第一行数据
data_df.describe() # 每列特征的简单统计信息

drop_features = ['blueGoldDiff', 'redGoldDiff', 
                 'blueExperienceDiff', 'redExperienceDiff', 
                 'blueCSPerMin', 'redCSPerMin', 
                 'blueGoldPerMin', 'redGoldPerMin'] # 需要舍去的特征列
df = data_df.drop(columns=drop_features) # 舍去特征列
info_names = [c[3:] for c in df.columns if c.startswith('red')] # 取出要作差值的特征名字（除去red前缀）
for info in info_names: # 对于每个特征名字
    df['br' + info] = df['blue' + info] - df['red' + info] # 构造一个新的特征，由蓝色特征减去红色特征，前缀为br
# 其中FirstBlood为首次击杀最多有一只队伍能获得，brFirstBlood=1为蓝，0为没有产生，-1为红
df = df.drop(columns=['blueFirstBlood', 'redFirstBlood']) # 原有的FirstBlood可删除

DISCRETE_N = 10
discrete_df = df.copy() # 先复制一份数据
for c in df.columns[1:]: # 遍历每一列特征，跳过标签列
    if len(df[c].unique()) <= DISCRETE_N: # 对原本取值可能性就较少的列，比如大龙数目，可不作处理
        continue
    else:
        # precsion=0表示区间间隔数字的精度保持和数据原本精度相同，duplicates='drop'表示去除某些重复的分隔数（即空区间）
        # 调用pandas的qcut函数，表示等密度区间，尽量让每个区间里的样本数相当，但每个区间长度可能不同
        discrete_df[c] = pd.qcut(df[c], DISCRETE_N, precision=0, labels=False, duplicates='drop')
        # 调用pandas的cut函数，表示等长的区间，尽量让每个区间的长度相当，但每个区间里的样本数可能不同
        # discrete_df[c] = pd.cut(df[c], 10, precision=0, duplicates='drop')

all_y = discrete_df['blueWins'].values # 所有标签数据
feature_names = discrete_df.columns[1:] # 所有特征的名称
all_x = discrete_df[feature_names].values # 所有原始特征值，pandas的DataFrame.values取出为numpy的array矩阵

# 划分训练集和测试集
x_train, x_test, y_train, y_test = train_test_split(all_x, all_y, test_size=0.2, random_state=RANDOM_SEED)
all_y.shape, all_x.shape, x_train.shape, x_test.shape, y_train.shape, y_test.shape # 输出数据行列信息

# 定义决策树类
class DecisionTree(object):
    def __init__(self, classes, features, 
                 max_depth=10, min_samples_split=10,
                 impurity_t='entropy'):
        '''
        传入一些可能用到的模型参数，也可能不会用到
        classes表示模型分类总共有几类
        features是每个特征的名字，也方便查询总的共特征数
        max_depth表示构建决策树时的最大深度
        min_samples_split表示构建决策树分裂节点时，如果到达该节点的样本数小于该值则不再分裂
        impurity_t表示计算混杂度（不纯度）的计算方式，例如entropy或gini
        '''  
        self.classes = classes
        self.features = features
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.impurity_t = impurity_t
        self.root = None # 定义根节点，未训练时为空
    
    def impurity(self, data):
        '''
        计算混杂度
        参数data是一个numpy一维数组
        '''
        cnt = Counter(data) # 计数每个值出现的次数
        ps = [1.0 * cnt[v]/len(data) for v in cnt] # 计算每个值出现的比例
        if self.impurity_t == 'entropy': # 如果是信息熵
            return -np.sum([p * np.log2(p) for p in ps if p > 0]), cnt # 同时返回每个值出现的次数可能方便之后运算
        return 1 - np.sum([p*p for p in ps]), cnt # 返回gini系数
    
    def gain(self, feature, label):
        '''
        计算某个特征下的信息增益
        feature是特征的值，numpy一维数组
        label是对应的标签，numpy一维数组
        '''
        c_impurity, _ = self.impurity(label) # 不考虑特征时标签的混杂度
        
        # 记录特征的每种取值所对应的数组下标
        f_index = {} 
        for idx, v in enumerate(feature):
            if v not in f_index:
                f_index[v] = []
            f_index[v].append(idx)
        
        # 计算根据该特征分裂后的不纯度，根据特征的每种值的数目加权和
        f_impurity = 0
        for v in f_index:
            f_l = label[f_index[v]] # 取出该特征取值对应的数组下标
            f_impurity += self.impurity(f_l)[0] * len(f_l)/len(label) # 计算不纯度并乘以该特征取值的比例
        
        
        # 有些特征取值很多，天然不纯度高、信息增益高，模型会偏向于取值很多的特征比如日期，但很可能过拟合
        # 计算信息增益率可以缓解该问题
        r = self.impurity(feature)[0] # 计算该特征在标签无关时的不纯度
        r = (c_impurity - f_impurity)/r if r > 0 else c_impurity - f_impurity # 除数不为0时为信息增益率
        return r, f_index # 返回信息增益率，以及每个特征取值的数组下标，方便之后使用
    
    
    def expand_node(self, feature, label, depth, skip_features=set()):
        '''
        分裂节点，feature和label为到达该节点的样本
        feature为二维numpy（n*m）数组，每行表示一个样本，有m个特征
        label为一维numpy（n）数组，表示每个样本的分类标签
        depth记录了当前节点的深度
        skip_features表示当前路径已经用到的特征
        在当前ID3算法离散特征的实现下，一条路径上已经用过的特征不会再用（其他实现有可能会选重复特征）
        '''
        l_cnt = Counter(label) # 计数每个类别的样本出现次数
        if len(l_cnt) == 0: # 如果只有一种类别了，无需分裂，已经是叶节点
            return label[0] # 只需记录类别
        if len(label) < self.min_samples_split or depth > self.max_depth: # 如果达到了最小分裂的样本数或者最大深度的阈值
            return l_cnt.most_common(1)[0][0] # 则只记录当前样本中最多的类别
        
        f_idx, max_gain, fv_index = -1, -1, None # 准备挑选分裂特征
        for idx in range(len(self.features)): # 遍历所有特征
            if idx in skip_features: # 如果当前路径已经用到，不用再算
                continue
            f_gain, fv = self.gain(feature[:, idx], label) # 计算特征的信息增益，fv是特征每个取值的样本下标
            
            # if f_gain <= 0: # 如果信息增益不为正，跳过该特征
            #    continue
            if f_idx < 0 or f_gain > max_gain: # 如果个更好的分裂特征
                f_idx, max_gain, fv_index = idx, f_gain, fv # 则记录该特征
        
        if f_idx < 0: # 如果没有找到合适的特征，即所有特征都没有信息增益
            return l_cnt.most_common(1)[0][0] # 则只记录当前样本中最多的类别
            
        decision = {} # 用字典记录每个特征取值所对应的子节点，key是特征取值，value是子节点
        skip_features = set([f_idx]+[f for f in skip_features]) # 子节点要跳过的特征包括当前选择的特征
        for v in fv_index: # 遍历特征的每种取值
            decision[v] = self.expand_node(feature[fv_index[v], :], label[fv_index[v]],  # 取出该特征取值所对应的样本
                                           depth=depth + 1, skip_features=skip_features) # 深度+1，递归调用节点分裂
        # 返回一个元组，有三个元素
        # 第一个是选择的特征下标，第二个特征取值和对应的子节点，第三个是到达当前节点的样本中最多的类别
        return (f_idx, decision, l_cnt.most_common(1)[0][0]) 
        
    def traverse_node(self, node, feature):
        '''
        预测样本时从根节点开始遍历节点，根据特征路由。
        node表示当前到达的节点，例如self.root
        feature是长度为m的numpy一维数组
        '''
        assert len(self.features) == len(feature) # 要求输入样本特征数和模型定义时特征数目一致
        if type(node) is not tuple: # 如果到达了一个节点是叶节点（不再分裂），则返回该节点类别
            return node
        fv = feature[node[0]] # 否则取出该节点对应的特征值，node[0]记录了特征的下标
        if fv in node[1]: # 根据特征值找到子节点，注意需要判断训练节点分裂时到达该节点的样本是否有该特征值（分支）
            return self.traverse_node(node[1][fv], feature) # 如果有，则进入到子节点继续遍历
        return node[-1] # 如果没有，返回训练时到达当前节点的样本中最多的类别
        
    def fit(self, feature, label):
        '''
        训练模型
        feature为二维numpy（n*m）数组，每行表示一个样本，有m个特征
        label为一维numpy（n）数组，表示每个样本的分类标签
        '''
        assert len(self.features) == len(feature[0]) # 输入数据的特征数目应该和模型定义时的特征数目相同
        self.root = self.expand_node(feature, label, depth=1) # 从根节点开始分裂，模型记录根节点
        
    
    def predict(self, feature):
        '''
        预测
        输入feature可以是一个一维numpy数组也可以是一个二维numpy数组
        如果是一维numpy（m）数组则是一个样本，包含m个特征，返回一个类别值
        如果是二维numpy（n*m）数组则表示n个样本，每个样本包含m个特征，返回一个numpy一维数组
        '''
        assert len(feature.shape) == 1 or len(feature.shape) == 2 # 只能是1维或2维
        if len(feature.shape) == 1: # 如果是一个样本
            return self.traverse_node(self.root, feature) # 从根节点开始路由
        return np.array([self.traverse_node(self.root, f) for f in feature]) # 如果是很多个样本
    
    def get_params(self, deep): # 要调用sklearn的cross_validate需要实现该函数返回所有参数
        return {'classes': self.classes, 'features': self.features, 
                'max_depth': self.max_depth, 'min_samples_split': self.min_samples_split,
                'impurity_t': self.impurity_t}

    def set_params(self, **parameters): # 要调用sklearn的GridSearchCV需要实现该函数给类设定所有参数
        for parameter, value in parameters.items():
            setattr(self, parameter, value)
        return self
        
        
# 定义决策树模型，传入算法参数
DT = DecisionTree(classes=[0,1], features=feature_names, max_depth=5, min_samples_split=10, impurity_t='gini')
# DT = DecisionTreeClassifier(max_depth=5, min_samples_split=10)

DT.fit(x_train, y_train) # 在训练集上训练
p_test = DT.predict(x_test) # 在测试集上预测，获得预测值
print(p_test) # 输出预测值
test_acc = accuracy_score(p_test, y_test) # 将测试预测值与测试集标签对比获得准确率
print('accuracy: {:.4f}'.format(test_acc)) # 输出准确率

best = None # 记录最佳结果
for impurity_t in ['entropy', 'gini']: # 遍历不纯度的计算方式
    for max_depth in range(1, 6): # 遍历最大树深度
        for min_samples_split in [50, 100, 200, 500, 1000]: # 遍历节点分裂最小样本数的阈值
            DT = DecisionTree(classes=[0,1], features=feature_names,  # 定义决策树
                              max_depth=max_depth, min_samples_split=min_samples_split, impurity_t=impurity_t)
            cv_result = cross_validate(DT, all_x, all_y, scoring=('accuracy'), cv=5) # 5折交叉验证            
            test_acc = np.mean(cv_result['test_score']) # 5折平均准确率
            current = (test_acc, max_depth, min_samples_split, impurity_t) # 记录参数和结果
            if best is None or test_acc > best[0]: # 如果是比当前最佳更好的结果
                best = current # 记录最好结果
                print('better accuracy: {:.4f}, max_depth={}, min_samples_split={}, impurity_t={}'.format(*best)) # 输出准确率和参数
            else:
                print('accuracy: {:.4f}, max_depth={}, min_samples_split={}, impurity_t={}'.format(*current)) # 输出准确率和参数

parameters = {'impurity_t':['entropy', 'gini'], 
              'max_depth': range(1, 6), 
              'min_samples_split': [50, 100, 200, 500, 1000]} # 定义需要遍历的参数
DT = DecisionTree(classes=[0,1], features=feature_names) # 定义决策树，可以不传参数，由GridSearchCV传入构建
grid_search = GridSearchCV(DT, parameters, scoring='accuracy', cv=5, verbose=10, n_jobs=4) # 传入模型和要遍历的参数
grid_search.fit(all_x, all_y) # 在所有数据上搜索参数
print(grid_search.best_score_, grid_search.best_params_) # 输出最佳指标和最佳参数

best_dt = grid_search.best_estimator_ # 取出最佳模型
print('root', best_dt.features[best_dt.root[0]]) # 输出根节点特征
for fv in best_dt.root[1]: # 遍历根节点的每种特征取值
    print(fv, '->', best_dt.features[best_dt.root[1][fv][0]], '-> ...') # 输出下一层特征