# src/grid_search_analyzer.py

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from typing import List, Dict, Any, Tuple

class GridSearchAnalyzer:
    """
    用于分析和可视化网格搜索结果的类。
    """
    def __init__(self, grid_results: List[Dict[str, Any]], optimization_metric: str):
        """
        初始化 GridSearchAnalyzer。

        Args:
            grid_results (List[Dict[str, Any]]): 包含每个参数组合及其对应优化指标的列表。
                                                 例如：[{'param1': v1, 'param2': v2, 'metric': m1}, ...]
            optimization_metric (str): 用于优化的指标名称（例如 'total_return'）。
        """
        if not grid_results:
            raise ValueError("grid_results 列表不能为空。")
        if optimization_metric not in grid_results[0]:
            raise ValueError(f"优化指标 '{optimization_metric}' 不在 grid_results 的字典中。")

        self.grid_results = grid_results
        self.optimization_metric = optimization_metric
        self.param_names = [k for k in grid_results[0].keys() if k != optimization_metric]

        if len(self.param_names) != 2:
            raise ValueError("GridSearchAnalyzer 当前只支持分析两个参数的网格搜索结果。")

        self.param1_name = self.param_names[0]
        self.param2_name = self.param_names[1]

    def find_best_parameters(self) -> Dict[str, Any]:
        """
        找到网格搜索中表现最佳的参数组合和其指标值。
        """
        if not self.grid_results:
            return {}
        best_result = max(self.grid_results, key=lambda x: x[self.optimization_metric])
        print(f"\n--- 最佳参数组合 ---")
        print(f"  {self.param1_name}: {best_result[self.param1_name]}")
        print(f"  {self.param2_name}: {best_result[self.param2_name]}")
        print(f"  {self.optimization_metric}: {best_result[self.optimization_metric]:.4f}")
        return best_result

    def plot_heatmap(self, title: str = "heatmap"):
        """
        绘制两个参数的热力图。
        """
        if not self.grid_results:
            print("没有数据用于绘制热力图。")
            return

        x_values = sorted(list(set(d[self.param1_name] for d in self.grid_results)))
        y_values = sorted(list(set(d[self.param2_name] for d in self.grid_results)))

        heatmap_matrix = np.zeros((len(y_values), len(x_values)))

        # 填充网格
        for item in self.grid_results:
            x_idx = x_values.index(item[self.param1_name])
            y_idx = y_values.index(item[self.param2_name])
            heatmap_matrix[y_idx, x_idx] = item[self.optimization_metric]

        print([x_values[0], x_values[-1], y_values[0], y_values[-1]])
        fig, ax = plt.subplots(figsize=(10, 8))
        im = ax.imshow(heatmap_matrix, cmap='viridis', origin='lower',
                       extent=[x_values[0], x_values[-1], y_values[0], y_values[-1]],
                       aspect='auto')

        ax.set_xticks(x_values)
        ax.set_yticks(y_values)

        ax.set_xlabel(self.param1_name)
        ax.set_ylabel(self.param2_name)
        ax.set_title(f"{title}: {self.optimization_metric} vs. {self.param1_name} & {self.param2_name}")

        cbar = fig.colorbar(im, ax=ax)
        cbar.set_label(self.optimization_metric)

        # 在每个格子上显示数值
        for i in range(len(y_values)):
            for j in range(len(x_values)):
                text = ax.text(x_values[j], y_values[i], f"{heatmap_matrix[i, j]:.2f}",
                               ha="center", va="center", color="w", fontsize=8)

        plt.tight_layout()
        plt.show()