import numpy as np
import matplotlib.pyplot as plt

class RegressionTreeNode:
    """回帰木のノード"""
    def __init__(self, value=None, feature_index=None, threshold=None, left=None, right=None):
        self.value = value              # リーフの予測値
        self.feature_index = feature_index  # 分割特徴量
        self.threshold = threshold      # 分割閾値
        self.left = left                # 左の子ノード
        self.right = right              # 右の子ノード


class RegressionTree:
    """回帰木（CART）"""
    def __init__(self, max_depth=3, min_samples_leaf=5):
        self.max_depth = max_depth
        self.min_samples_leaf = min_samples_leaf
        self.root = None

    def _mse(self, y):
        """平均二乗誤差"""
        if len(y) == 0:
            return 0
        return np.var(y) * len(y)

    def _best_split(self, X, y):
        """最適な分割を探索"""
        n_samples, n_features = X.shape
        best_gain = -float('inf')
        best_feature = None
        best_threshold = None

        current_mse = self._mse(y)

        for feature_idx in range(n_features):
            thresholds = np.unique(X[:, feature_idx])
            for threshold in thresholds:
                left_mask = X[:, feature_idx] <= threshold
                right_mask = ~left_mask

                if np.sum(left_mask) < self.min_samples_leaf or np.sum(right_mask) < self.min_samples_leaf:
                    continue

                gain = current_mse - self._mse(y[left_mask]) - self._mse(y[right_mask])
                if gain > best_gain:
                    best_gain = gain
                    best_feature = feature_idx
                    best_threshold = threshold

        return best_feature, best_threshold, best_gain

    def _build_tree(self, X, y, depth):
        """再帰的に木を構築"""
        # リーフノードの条件
        if depth >= self.max_depth or len(y) < 2 * self.min_samples_leaf:
            return RegressionTreeNode(value=np.mean(y))

        feature_idx, threshold, gain = self._best_split(X, y)
        if feature_idx is None or gain <= 0:
            return RegressionTreeNode(value=np.mean(y))

        left_mask = X[:, feature_idx] <= threshold
        right_mask = ~left_mask

        left_child = self._build_tree(X[left_mask], y[left_mask], depth + 1)
        right_child = self._build_tree(X[right_mask], y[right_mask], depth + 1)

        return RegressionTreeNode(
            feature_index=feature_idx,
            threshold=threshold,
            left=left_child,
            right=right_child
        )

    def fit(self, X, y):
        """学習"""
        self.root = self._build_tree(X, y, depth=0)

    def _predict_single(self, x, node):
        """1サンプルの予測"""
        if node.value is not None:
            return node.value
        if x[node.feature_index] <= node.threshold:
            return self._predict_single(x, node.left)
        else:
            return self._predict_single(x, node.right)

    def predict(self, X):
        """予測"""
        return np.array([self._predict_single(x, self.root) for x in X])


class GradientBoostingRegressor:
    """勾配ブースティング回帰のスクラッチ実装"""
    def __init__(self, n_estimators=100, learning_rate=0.1, max_depth=3, min_samples_leaf=5):
        self.n_estimators = n_estimators
        self.learning_rate = learning_rate
        self.max_depth = max_depth
        self.min_samples_leaf = min_samples_leaf
        self.trees = []
        self.f0 = None  # 初期予測値
        self.train_losses = []

    def fit(self, X, y):
        """学習"""
        # 初期値: ターゲットの平均
        self.f0 = np.mean(y)
        f = np.full(len(y), self.f0)

        for t in range(self.n_estimators):
            # 擬似残差の計算（二乗誤差損失の場合は残差そのもの）
            residuals = y - f

            # 回帰木を擬似残差にフィット
            tree = RegressionTree(max_depth=self.max_depth, min_samples_leaf=self.min_samples_leaf)
            tree.fit(X, residuals)

            # モデルの更新
            update = tree.predict(X)
            f += self.learning_rate * update

            self.trees.append(tree)

            # 訓練損失の記録
            loss = np.mean((y - f) ** 2)
            self.train_losses.append(loss)

    def predict(self, X):
        """予測"""
        f = np.full(X.shape[0], self.f0)
        for tree in self.trees:
            f += self.learning_rate * tree.predict(X)
        return f

# 合成データの生成
np.random.seed(42)
X_train = np.sort(np.random.uniform(0, 10, 200)).reshape(-1, 1)
y_true = np.sin(X_train.ravel()) + 0.5 * np.sin(3 * X_train.ravel())
y_train = y_true + np.random.normal(0, 0.3, len(y_true))

# 勾配ブースティングの学習
gb = GradientBoostingRegressor(n_estimators=200, learning_rate=0.1, max_depth=3, min_samples_leaf=5)
gb.fit(X_train, y_train)

# --- 段階的な予測の可視化 ---
fig, axes = plt.subplots(2, 3, figsize=(16, 10))
X_test = np.linspace(0, 10, 500).reshape(-1, 1)

rounds_to_show = [1, 5, 10, 30, 100, 200]

for idx, n_trees in enumerate(rounds_to_show):
    ax = axes[idx // 3, idx % 3]

    # n_trees ラウンドまでの累積予測
    f_pred = np.full(X_test.shape[0], gb.f0)
    for t in range(min(n_trees, len(gb.trees))):
        f_pred += gb.learning_rate * gb.trees[t].predict(X_test)

    ax.scatter(X_train, y_train, c='gray', s=10, alpha=0.5, label='Training data')
    ax.plot(X_test, np.sin(X_test.ravel()) + 0.5 * np.sin(3 * X_test.ravel()),
            'g--', linewidth=2, label='True function')
    ax.plot(X_test, f_pred, 'r-', linewidth=2, label='GBDT prediction')
    ax.set_title(f'T = {min(n_trees, len(gb.trees))}', fontsize=13)
    ax.legend(fontsize=9)
    ax.set_ylim(-2.5, 2.5)
    ax.grid(True, alpha=0.3)

plt.suptitle('Gradient Boosting Regression: Prediction at Different Rounds', fontsize=15, y=1.02)
plt.tight_layout()
plt.savefig('gbdt_regression_progress.png', dpi=150, bbox_inches='tight')
plt.show()

# --- 訓練損失の推移 ---
fig, ax = plt.subplots(figsize=(10, 5))
ax.plot(range(1, len(gb.train_losses) + 1), gb.train_losses, 'b-', linewidth=2)
ax.set_xlabel('Number of estimators', fontsize=12)
ax.set_ylabel('Training MSE', fontsize=12)
ax.set_title('Training Loss vs Number of Estimators', fontsize=14)
ax.set_yscale('log')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('gbdt_training_loss.png', dpi=150, bbox_inches='tight')
plt.show()

from sklearn.ensemble import GradientBoostingRegressor as SklearnGBR
from sklearn.metrics import mean_squared_error

# scikit-learn のGBDT
sklearn_gb = SklearnGBR(n_estimators=200, learning_rate=0.1, max_depth=3, random_state=42)
sklearn_gb.fit(X_train, y_train)

# テストデータでの予測
X_test = np.linspace(0, 10, 500).reshape(-1, 1)
y_test_true = np.sin(X_test.ravel()) + 0.5 * np.sin(3 * X_test.ravel())

y_pred_scratch = gb.predict(X_test)
y_pred_sklearn = sklearn_gb.predict(X_test)

# 比較の可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

axes[0].scatter(X_train, y_train, c='gray', s=10, alpha=0.5, label='Training data')
axes[0].plot(X_test, y_test_true, 'g--', linewidth=2, label='True function')
axes[0].plot(X_test, y_pred_scratch, 'r-', linewidth=2, label='Scratch GBDT')
axes[0].set_title('Scratch Implementation', fontsize=14)
axes[0].legend(fontsize=10)
axes[0].grid(True, alpha=0.3)

axes[1].scatter(X_train, y_train, c='gray', s=10, alpha=0.5, label='Training data')
axes[1].plot(X_test, y_test_true, 'g--', linewidth=2, label='True function')
axes[1].plot(X_test, y_pred_sklearn, 'b-', linewidth=2, label='scikit-learn GBDT')
axes[1].set_title('scikit-learn Implementation', fontsize=14)
axes[1].legend(fontsize=10)
axes[1].grid(True, alpha=0.3)

plt.suptitle('Gradient Boosting: Scratch vs scikit-learn', fontsize=15, y=1.02)
plt.tight_layout()
plt.savefig('gbdt_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

# MSEの比較
mse_scratch = mean_squared_error(y_test_true, y_pred_scratch)
mse_sklearn = mean_squared_error(y_test_true, y_pred_sklearn)
print(f"Scratch GBDT - MSE: {mse_scratch:.4f}")
print(f"scikit-learn GBDT - MSE: {mse_sklearn:.4f}")

# 学習率の異なるGBDTの比較
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
learning_rates = [0.01, 0.1, 0.5]

for idx, lr in enumerate(learning_rates):
    ax = axes[idx]
    gb_lr = GradientBoostingRegressor(n_estimators=200, learning_rate=lr, max_depth=3)
    gb_lr.fit(X_train, y_train)

    y_pred = gb_lr.predict(X_test)
    ax.scatter(X_train, y_train, c='gray', s=10, alpha=0.5)
    ax.plot(X_test, y_test_true, 'g--', linewidth=2, label='True')
    ax.plot(X_test, y_pred, 'r-', linewidth=2, label=f'GBDT (lr={lr})')
    ax.set_title(f'Learning Rate = {lr}', fontsize=13)
    ax.legend(fontsize=10)
    ax.set_ylim(-2.5, 2.5)
    ax.grid(True, alpha=0.3)

plt.suptitle('Effect of Learning Rate on GBDT', fontsize=15, y=1.02)
plt.tight_layout()
plt.savefig('gbdt_learning_rate.png', dpi=150, bbox_inches='tight')
plt.show()