import numpy as np
import matplotlib.pyplot as plt

import tkml_physbo as wrapper


def main():
    # 1. 設定
    infile = "result.xlsx"  # 入力ファイル
    target_var_idx = 0      # 感度解析を行いたい記述子のインデックス
    
    # 2. データのロードと前処理
    print("--- Loading and Preprocessing ---")
    df = wrapper.load_data(infile)
    X_raw, y_raw, feat_names = wrapper.preprocess(df)
    
    # 3. 標準化
    X_std, y_std, s_dict = wrapper.transform(X_raw.values, y_raw)
    
    # 4. モデルの構築と学習
    print("--- Training PHYSBO Model ---")
    # 設定のカスタマイズ（基底数や探索回数など）
    config = wrapper.model({
        "num_rand_basis": 500, 
        "max_num_probes": 0,    # 回帰のみの場合は0
        "score_mode": 'EI'
    })
    trained_policy = wrapper.fit(X_std, y_std, X_std, config)
    
    # 5. 予測と評価
    y_pred_std, y_var_std = wrapper.predict(trained_policy, X_std)
    y_pred_orig = wrapper.inverse_transform(y_pred_std, s_dict)
    
    # 評価指標の表示
    res_metrics = wrapper.metrics(y_raw, y_pred_orig)
    print("\nModel Metrics:")
    for k, v in res_metrics.items():
        print(f"  {k}: {v:.4f}")

    # 6. 可視化：YY-plot (実測値 vs 予測値)
    plt.figure(figsize=(6, 5))
    plt.scatter(y_raw, y_pred_orig, alpha=0.7, edgecolors='k')
    
    # 45度線（理想線）
    min_val = min(y_raw.min(), y_pred_orig.min())
    max_val = max(y_raw.max(), y_pred_orig.max())
    plt.plot([min_val, max_val], [min_val, max_val], 'r--', label='Ideal')
    
    plt.xlabel("Actual Value")
    plt.ylabel("Predicted Value")
    plt.title(f"YY-Plot (R2: {res_metrics['R2']:.3f})")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("yy_plot.png")
    plt.show()

    # 7. 可視化：感度解析 (Sensitivity Analysis)
    print(f"\n--- Sensitivity Analysis for: {feat_names[target_var_idx]} ---")
    # 変化させる範囲を元のスケールで定義（例：データの最小から最大まで）
    v_min = X_raw.iloc[:, target_var_idx].min()
    v_max = X_raw.iloc[:, target_var_idx].max()
    var_range = np.linspace(v_min, v_max, 50)
    
    vals, responses = wrapper.sensitivity_analysis(
        trained_policy, s_dict, feat_names, target_var_idx, var_range
    )
    
    plt.figure(figsize=(8, 5))
    plt.plot(vals, responses, label='Prediction', color='blue', lw=2)
    plt.xlabel(feat_names[target_var_idx])
    plt.ylabel("Predicted Objective Value")
    plt.title(f"Sensitivity Analysis: {feat_names[target_var_idx]}")
    plt.grid(True, linestyle=':')
    plt.tight_layout()
    plt.savefig("sensitivity_analysis.png")
    plt.show()

    # 8. 結果の保存
    wrapper.save_data(X_raw, y_pred_orig, "result_with_predictions.xlsx", 
                      descriptor_names=feat_names, objective_name="o:predicted_mean")
    print("\nAnalysis complete. Results saved.")

if __name__ == "__main__":
    main()