import numpy as np
import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
from sklearn.datasets import make_classification, make_regression
from sklearn.ensemble import RandomForestClassifier as SklearnRF
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, mean_squared_error

# ----------------------------
# ランダムフォレスト分類器（スクラッチ実装）
# ----------------------------
class RandomForestClassifierScratch:
    """ランダムフォレスト分類器（DecisionTreeClassifierベース）"""

    def __init__(self, n_estimators=100, max_features='sqrt',
                 max_depth=None, random_state=None):
        self.n_estimators = n_estimators
        self.max_features = max_features
        self.max_depth = max_depth
        self.random_state = random_state
        self.estimators = []
        self.feature_indices = []  # 各木で使用する特徴量インデックス
        self.oob_indices = []
        self.feature_importances_mdi = None

    def _get_max_features(self, n_features):
        """max_featuresの数値を取得"""
        if self.max_features == 'sqrt':
            return max(1, int(np.sqrt(n_features)))
        elif self.max_features == 'log2':
            return max(1, int(np.log2(n_features)))
        elif isinstance(self.max_features, int):
            return self.max_features
        elif isinstance(self.max_features, float):
            return max(1, int(self.max_features * n_features))
        else:
            return n_features

    def fit(self, X, y):
        """ランダムフォレストの学習"""
        n_samples, n_features = X.shape
        self.n_features_ = n_features
        self.classes_ = np.unique(y)
        m = self._get_max_features(n_features)
        rng = np.random.RandomState(self.random_state)

        self.estimators = []
        self.feature_indices = []
        self.oob_indices = []

        # MDI重要度の蓄積
        importances = np.zeros(n_features)

        for b in range(self.n_estimators):
            # ブートストラップサンプリング
            bootstrap_idx = rng.choice(n_samples, size=n_samples, replace=True)
            oob_idx = np.setdiff1d(np.arange(n_samples), np.unique(bootstrap_idx))
            self.oob_indices.append(oob_idx)

            X_boot = X[bootstrap_idx]
            y_boot = y[bootstrap_idx]

            # 決定木の学習（max_features で特徴量をランダム選択）
            tree = DecisionTreeClassifier(
                max_depth=self.max_depth,
                max_features=m,
                random_state=rng.randint(0, 100000)
            )
            tree.fit(X_boot, y_boot)
            self.estimators.append(tree)

            # MDI重要度の加算
            importances += tree.feature_importances_

        # MDI重要度の正規化
        self.feature_importances_mdi = importances / self.n_estimators

        return self

    def predict(self, X):
        """多数決による予測"""
        predictions = np.array([tree.predict(X) for tree in self.estimators])
        # 各サンプルで最も多いクラスを選択
        result = np.zeros(X.shape[0], dtype=self.classes_.dtype)
        for i in range(X.shape[0]):
            values, counts = np.unique(predictions[:, i], return_counts=True)
            result[i] = values[np.argmax(counts)]
        return result

    def predict_proba(self, X):
        """クラス確率の推定"""
        n_classes = len(self.classes_)
        proba = np.zeros((X.shape[0], n_classes))
        for tree in self.estimators:
            tree_proba = tree.predict_proba(X)
            # クラスの対応付け
            for j, c in enumerate(tree.classes_):
                idx = np.where(self.classes_ == c)[0][0]
                proba[:, idx] += tree_proba[:, j]
        proba /= self.n_estimators
        return proba

    def oob_score(self, X, y):
        """OOBスコア（正確度）"""
        n_samples = X.shape[0]
        n_classes = len(self.classes_)
        oob_votes = np.zeros((n_samples, n_classes))

        for b, tree in enumerate(self.estimators):
            oob_idx = self.oob_indices[b]
            if len(oob_idx) > 0:
                tree_proba = tree.predict_proba(X[oob_idx])
                for j, c in enumerate(tree.classes_):
                    idx = np.where(self.classes_ == c)[0][0]
                    oob_votes[oob_idx, idx] += tree_proba[:, j]

        valid = np.sum(oob_votes, axis=1) > 0
        oob_pred = self.classes_[np.argmax(oob_votes[valid], axis=1)]
        return np.mean(oob_pred == y[valid])

    def permutation_importance(self, X, y):
        """MDA（Permutation重要度）の計算"""
        base_score = self.oob_score(X, y)
        importances = np.zeros(self.n_features_)

        for j in range(self.n_features_):
            X_permuted = X.copy()
            perm_idx = np.random.permutation(X.shape[0])
            X_permuted[:, j] = X[perm_idx, j]
            perm_score = self.oob_score(X_permuted, y)
            importances[j] = base_score - perm_score

        return importances

    def proximity_matrix(self, X):
        """Proximity行列の計算"""
        n_samples = X.shape[0]
        prox = np.zeros((n_samples, n_samples))

        for tree in self.estimators:
            leaves = tree.apply(X)
            for i in range(n_samples):
                for j in range(i, n_samples):
                    if leaves[i] == leaves[j]:
                        prox[i, j] += 1
                        prox[j, i] += 1

        prox /= self.n_estimators
        return prox

# ----------------------------
# データ生成
# ----------------------------
np.random.seed(42)

# 分類タスク
X_cls, y_cls = make_classification(
    n_samples=500, n_features=20, n_informative=10,
    n_redundant=5, n_classes=2, random_state=42
)
X_train_cls, X_test_cls, y_train_cls, y_test_cls = train_test_split(
    X_cls, y_cls, test_size=0.3, random_state=42
)

# ----------------------------
# 実験1: スクラッチ実装 vs scikit-learn の比較
# ----------------------------
print("=== スクラッチ実装 vs scikit-learn ===")

# スクラッチ実装
rf_scratch = RandomForestClassifierScratch(
    n_estimators=100, max_features='sqrt', random_state=42
)
rf_scratch.fit(X_train_cls, y_train_cls)
y_pred_scratch = rf_scratch.predict(X_test_cls)
acc_scratch = accuracy_score(y_test_cls, y_pred_scratch)

# scikit-learn
rf_sklearn = SklearnRF(
    n_estimators=100, max_features='sqrt', random_state=42
)
rf_sklearn.fit(X_train_cls, y_train_cls)
y_pred_sklearn = rf_sklearn.predict(X_test_cls)
acc_sklearn = accuracy_score(y_test_cls, y_pred_sklearn)

print(f"スクラッチ実装  正確度: {acc_scratch:.4f}")
print(f"scikit-learn   正確度: {acc_sklearn:.4f}")

# OOBスコア
oob_scratch = rf_scratch.oob_score(X_train_cls, y_train_cls)
print(f"\nスクラッチ OOBスコア: {oob_scratch:.4f}")
print(f"scikit-learn OOBスコア: {rf_sklearn.oob_score_:.4f}"
      if hasattr(rf_sklearn, 'oob_score_') else "")

# ----------------------------
# 実験2: 特徴量数mの影響（相関 vs 精度のトレードオフ）
# ----------------------------
print("\n=== 特徴量数 m の影響 ===")
m_values = [1, 2, 3, 4, 5, 7, 10, 15, 20]
accuracies = []
oob_scores_list = []

for m in m_values:
    rf = RandomForestClassifierScratch(
        n_estimators=100, max_features=m, random_state=42
    )
    rf.fit(X_train_cls, y_train_cls)
    acc = accuracy_score(y_test_cls, rf.predict(X_test_cls))
    oob = rf.oob_score(X_train_cls, y_train_cls)
    accuracies.append(acc)
    oob_scores_list.append(oob)
    print(f"m = {m:2d}  Test Acc: {acc:.4f}  OOB Score: {oob:.4f}")

fig, ax = plt.subplots(figsize=(10, 6))
ax.plot(m_values, accuracies, 'bo-', linewidth=2, markersize=8, label='Test Accuracy')
ax.plot(m_values, oob_scores_list, 'rs--', linewidth=2, markersize=8, label='OOB Score')
ax.axvline(x=int(np.sqrt(20)), color='green', linestyle=':', linewidth=2,
           label=f'$\\sqrt{{p}}$ = {int(np.sqrt(20))}')
ax.set_xlabel('Number of Features (m)', fontsize=12)
ax.set_ylabel('Score', fontsize=12)
ax.set_title('Effect of max_features on Random Forest Performance', fontsize=14)
ax.legend(fontsize=11)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

# ----------------------------
# 実験3: 変数重要度の可視化
# ----------------------------
# MDI重要度
mdi_importance = rf_scratch.feature_importances_mdi

# scikit-learnの重要度と比較
sklearn_importance = rf_sklearn.feature_importances_

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))

# MDI重要度（スクラッチ）
sorted_idx_scratch = np.argsort(mdi_importance)
ax1.barh(range(len(mdi_importance)), mdi_importance[sorted_idx_scratch], color='steelblue')
ax1.set_yticks(range(len(mdi_importance)))
ax1.set_yticklabels([f'Feature {i}' for i in sorted_idx_scratch])
ax1.set_xlabel('MDI Importance')
ax1.set_title('Feature Importance (Scratch Implementation)')
ax1.grid(True, alpha=0.3, axis='x')

# scikit-learnの重要度
sorted_idx_sklearn = np.argsort(sklearn_importance)
ax2.barh(range(len(sklearn_importance)), sklearn_importance[sorted_idx_sklearn], color='coral')
ax2.set_yticks(range(len(sklearn_importance)))
ax2.set_yticklabels([f'Feature {i}' for i in sorted_idx_sklearn])
ax2.set_xlabel('MDI Importance')
ax2.set_title('Feature Importance (scikit-learn)')
ax2.grid(True, alpha=0.3, axis='x')

plt.tight_layout()
plt.show()

# ----------------------------
# 実験4: 推定器数 B と性能の関係
# ----------------------------
B_values = [1, 5, 10, 25, 50, 100, 200, 300]
test_accs = []
oob_accs = []

for B in B_values:
    rf = RandomForestClassifierScratch(
        n_estimators=B, max_features='sqrt', random_state=42
    )
    rf.fit(X_train_cls, y_train_cls)
    test_accs.append(accuracy_score(y_test_cls, rf.predict(X_test_cls)))
    oob_accs.append(rf.oob_score(X_train_cls, y_train_cls))

fig, ax = plt.subplots(figsize=(10, 6))
ax.plot(B_values, test_accs, 'bo-', linewidth=2, markersize=8, label='Test Accuracy')
ax.plot(B_values, oob_accs, 'rs--', linewidth=2, markersize=8, label='OOB Score')
ax.set_xlabel('Number of Trees (B)', fontsize=12)
ax.set_ylabel('Accuracy', fontsize=12)
ax.set_title('Random Forest: Accuracy vs Number of Trees', fontsize=14)
ax.legend(fontsize=11)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

# ----------------------------
# 実験5: Proximity行列の可視化（小規模データ）
# ----------------------------
# 小規模データで Proximity を計算
X_small, y_small = make_classification(
    n_samples=50, n_features=5, n_informative=3,
    n_classes=2, random_state=42
)

rf_small = RandomForestClassifierScratch(
    n_estimators=100, max_features='sqrt', random_state=42
)
rf_small.fit(X_small, y_small)

prox = rf_small.proximity_matrix(X_small)

# ソートしてクラスごとにまとめる
sort_idx = np.argsort(y_small)
prox_sorted = prox[sort_idx][:, sort_idx]

fig, ax = plt.subplots(figsize=(8, 7))
im = ax.imshow(prox_sorted, cmap='YlOrRd', vmin=0, vmax=1)
plt.colorbar(im, ax=ax, label='Proximity')
ax.set_xlabel('Sample Index (sorted by class)')
ax.set_ylabel('Sample Index (sorted by class)')
ax.set_title('Proximity Matrix (sorted by class label)')

# クラス境界線
n_class0 = np.sum(y_small == 0)
ax.axhline(y=n_class0 - 0.5, color='blue', linewidth=2)
ax.axvline(x=n_class0 - 0.5, color='blue', linewidth=2)

plt.tight_layout()
plt.show()

# 異常値スコアの計算
print("\n=== 異常値スコア（上位5件）===")
for k in np.unique(y_small):
    class_mask = y_small == k
    class_indices = np.where(class_mask)[0]
    n_k = len(class_indices)

    outlier_scores = np.zeros(n_k)
    for idx, i in enumerate(class_indices):
        prox_sum = np.sum(prox[i, class_indices] ** 2)
        if prox_sum > 0:
            outlier_scores[idx] = n_k / prox_sum

    top_outliers = np.argsort(outlier_scores)[-3:][::-1]
    print(f"\nクラス {k} の異常値スコア上位:")
    for rank, oi in enumerate(top_outliers):
        print(f"  {rank+1}. サンプル {class_indices[oi]}: "
              f"スコア = {outlier_scores[oi]:.4f}")