import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# データの読み込み
data = fetch_california_housing()
X = pd.DataFrame(data.data, columns=data.feature_names)
y = data.target

print(f"データサイズ: {X.shape}")
print(f"特徴量: {list(X.columns)}")
print(f"目的変数の統計量:\n{pd.Series(y).describe()}")

# GPRは計算コストがO(n^3)なのでサブサンプリング
n_subset = 1000
np.random.seed(42)
idx = np.random.choice(len(X), n_subset, replace=False)
X_sub = X.iloc[idx].values
y_sub = y[idx]

# 訓練/テストの分割
X_train, X_test, y_train, y_test = train_test_split(
    X_sub, y_sub, test_size=0.2, random_state=42
)

# 標準化（GPRではスケーリングが重要）
scaler_X = StandardScaler()
X_train_scaled = scaler_X.fit_transform(X_train)
X_test_scaled = scaler_X.transform(X_test)

scaler_y = StandardScaler()
y_train_scaled = scaler_y.fit_transform(y_train.reshape(-1, 1)).ravel()

print(f"\n訓練データ: {X_train_scaled.shape}")
print(f"テストデータ: {X_test_scaled.shape}")

from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, WhiteKernel, ConstantKernel, Matern
from sklearn.metrics import mean_squared_error, r2_score

# --- カーネルの定義 ---
# RBF + ノイズカーネル
kernel_rbf = ConstantKernel(1.0) * RBF(length_scale=1.0) + WhiteKernel(noise_level=0.1)

# Maternカーネル + ノイズ
kernel_matern = ConstantKernel(1.0) * Matern(length_scale=1.0, nu=2.5) + WhiteKernel(noise_level=0.1)

# --- 各カーネルでGPRを学習 ---
results = {}
kernels = {'RBF': kernel_rbf, 'Matern': kernel_matern}

for name, kernel in kernels.items():
    gpr = GaussianProcessRegressor(
        kernel=kernel,
        n_restarts_optimizer=5,
        random_state=42,
        normalize_y=False
    )
    gpr.fit(X_train_scaled, y_train_scaled)

    # 予測（平均と標準偏差）
    y_pred_scaled, y_std_scaled = gpr.predict(X_test_scaled, return_std=True)

    # 元のスケールに戻す
    y_pred = scaler_y.inverse_transform(y_pred_scaled.reshape(-1, 1)).ravel()
    y_std = y_std_scaled * scaler_y.scale_[0]

    rmse = np.sqrt(mean_squared_error(y_test, y_pred))
    r2 = r2_score(y_test, y_pred)

    results[name] = {
        'model': gpr,
        'y_pred': y_pred,
        'y_std': y_std,
        'rmse': rmse,
        'r2': r2
    }

    print(f"\n=== {name} カーネル ===")
    print(f"最適化後のカーネル: {gpr.kernel_}")
    print(f"対数周辺尤度: {gpr.log_marginal_likelihood_value_:.3f}")
    print(f"RMSE: {rmse:.4f}")
    print(f"R^2:  {r2:.4f}")

fig, axes = plt.subplots(2, 2, figsize=(14, 12))

for i, (name, res) in enumerate(results.items()):
    # 予測 vs 実測
    ax1 = axes[i, 0]
    ax1.scatter(y_test, res['y_pred'], alpha=0.5, s=20, c='blue')
    ax1.plot([y_test.min(), y_test.max()],
             [y_test.min(), y_test.max()], 'r--', linewidth=2)
    ax1.set_xlabel('Actual Price')
    ax1.set_ylabel('Predicted Price')
    ax1.set_title(f'{name}: Predicted vs Actual (R$^2$={res["r2"]:.3f})')
    ax1.grid(True, alpha=0.3)

    # 予測の不確実性
    ax2 = axes[i, 1]
    sorted_idx = np.argsort(y_test)
    ax2.errorbar(range(len(y_test)), res['y_pred'][sorted_idx],
                 yerr=2 * res['y_std'][sorted_idx],
                 fmt='none', alpha=0.3, color='blue')
    ax2.scatter(range(len(y_test)), y_test[sorted_idx],
                c='red', s=10, alpha=0.7, label='Actual')
    ax2.scatter(range(len(y_test)), res['y_pred'][sorted_idx],
                c='blue', s=10, alpha=0.7, label='Predicted')
    ax2.set_xlabel('Sample (sorted by actual)')
    ax2.set_ylabel('Price')
    ax2.set_title(f'{name}: Predictions with 95% CI')
    ax2.legend(fontsize=8)
    ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# --- 特徴量の重要度（長さスケール）---
for name, res in results.items():
    model = res['model']
    kernel = model.kernel_
    print(f"\n{name} カーネルのパラメータ:")
    for param_name, value in kernel.get_params().items():
        if 'length_scale' in param_name and isinstance(value, (float, np.floating, np.ndarray)):
            print(f"  {param_name}: {value}")