import numpy as np
import pandas as pd
from sklearn.datasets import load_boston
boston = load_boston()

X, y = boston.data, boston.target

y = y.reshape(-1, 1)
X = np.hstack((X, y))

def node_mean(X):
 return np.mean(X[:, -1])


def node_variance(X):
 return np.var(X[:, -1]) * X.shape[0]

from collections import defaultdict
def split_dataset(X, idx, thred):
 split_data = defaultdict(list)
 for x in X:
  split_data[x[idx] < thred].append(x)
 return list(split_data.values()), list(split_data.keys())

def get_thresholds(X, i):
 return set(X[:, i].tolist())

# 每次迭代方差优化的底线
MINIMUM_IMPROVE = 2.0
# 每个叶子节点最少样本数
MINIMUM_SAMPLES = 10

def split_variance(dataset, idx, threshold):
 left, right = [], []
 n = dataset.shape[0]
 for data in dataset:
  if data[idx] < threshold:
   left.append(data)
  else:
   right.append(data)
 left, right = np.array(left), np.array(right)
 # 预剪枝
 # 如果拆分结果有一边过少，则返回None，防止过拟合
 if len(left) < MINIMUM_SAMPLES or len(right) < MINIMUM_SAMPLES:
  return None
 # 拆分之后的方差和等于左子树的方差和加上右子树的方差和
 # 因为是方差和而不是均方差，所以可以累加
 return node_variance(left) + node_variance(right)

def choose_feature_to_split(dataset):
 n = len(dataset[0])-1
 m = len(dataset)
 # 记录最佳方差,特征和阈值
 var_ = node_variance(dataset)
 bestVar = float('inf')
 feature = -1
 thred = None
 for i in range(n):
  threds = get_thresholds(dataset, i)
  for t in threds:
   # 遍历所有的阈值，计算每个阈值的variance
   v = split_variance(dataset, i, t)
   # 如果v等于None，说明拆分过拟合了，跳过
   if v is None:
    continue
   if v < bestVar:
    bestVar, feature, thred = v, i, t
 # 如果最好的拆分效果达不到要求，那么就不拆分，控制子树的数量
 if var_ - bestVar < MINIMUM_IMPROVE:
  return None, None
 return feature, thred

def create_decision_tree(dataset):
 dataset = np.array(dataset)
 
 # 如果当前数量小于10，那么就不再继续划分了
 if dataset.shape[0] < MINIMUM_SAMPLES:
  return node_mean(dataset)
 
 # 记录最佳拆分的特征和阈值
 fidx, th = choose_feature_to_split(dataset)
 if fidx is None:
  return th
 
 node = {}
 node['feature'] = fidx
 node['threshold'] = th
 
 # 递归建树
 split_data, vals = split_dataset(dataset, fidx, th)
 for data, val in zip(split_data, vals):
  node[val] = create_decision_tree(data)
 return node

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=23)

def classify(node, data):
 key = node['feature']
 pred = None
 thred = node['threshold']

 if isinstance(node[data[key] < thred], dict):
  pred = classify(node[data[key] < thred], data)
 else:
  pred = node[data[key] < thred]
   
 # 放置pred为空，挑选一个叶子节点作为替补
 if pred is None:
  for key in node:
   if not isinstance(node[key], dict):
    pred = node[key]
    break
 return pred

def predict(node, X):
 y_pred = []
 for x in X:
  y = classify(node, x)
  y_pred.append(y)
 return np.array(y_pred)

def is_dict(node):
 return isinstance(node, dict)


def prune(node, testData):
 testData = np.array(testData)
 if testData.shape[0] == 0:
  return node
 
 # 拆分数据
 split_data, _ = split_dataset(testData, node['feature'], node['threshold'])
 # 对左右子树递归修剪
 if is_dict(node[0]):
  node[0] = prune(node[0], split_data[0])
 if is_dict(node[1]) and len(split_data) > 1:
  node[1] = prune(node[1], split_data[1])

 # 如果左右都是叶子节点，那么判断当前子树是否需要修剪
 if len(split_data) > 1 and not is_dict(node[0]) and not is_dict(node[1]):
  # 计算修剪前的方差和
  baseError = np.sum(np.power(np.array(split_data[0])[:, -1] - node[0], 2)) + np.sum(np.power(np.array(split_data[1])[:, -1] - node[1], 2))
  # 计算修剪后的方差和
  meanVal = (node[0] + node[1]) / 2
  mergeError = np.sum(np.power(meanVal - testData[:, -1], 2))
  if mergeError < baseError:
   return meanVal
  else:
   return node
 return node