import numpy as np
import pandas as pd
from scipy.special import logit as sp_logit
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import log_loss


def compute_metrics(
    df: pd.DataFrame,
    n_bins: int = 10,
    bin_strategy: str = "uniform",  # 'uniform' or 'quantile'
    include_draws: bool = True,
    eps: float = 1e-6,
) -> dict:
    """
    计算预测评估指标并拟合校准关系。

        参数:
            - df: 包含至少两列: 'win_prob' (预测主胜概率), 'res' (取 'won','refunded','lost')
      - n_bins: ECE 分箱数
      - bin_strategy: 'uniform' (等宽) 或 'quantile' (等频)
      - include_draws: 若 True, 将 'draw' 视为非胜 (y=0)。若 False, 丢弃 'draw' 行。
      - eps: 概率裁剪下限，用于数值稳定

    返回:
      dict 包含 logloss, brier, ece, accuracy, reg_alpha, reg_beta, ece_bins, n_samples
    """
    # 处理 refunded
    if include_draws:
        mask = df["res"].isin(["won", "refunded", "lost"])
    else:
        mask = df["res"].isin(["won", "lost"])
    df = df[mask].copy()

    # 标签: won=1, others=0 (包括 refunded)
    y = df["res"].map({"won": 1, "refunded": 0, "lost": 0}).astype(int).values
    p = df["win_prob"].astype(float).values

    # 裁剪概率以保证数值稳定
    p_clip = np.clip(p, eps, 1 - eps)

    # logloss: 使用 sklearn 实现以获得更稳健的数值行为
    try:
        logloss = float(log_loss(y, p_clip, labels=[0, 1]))
    except Exception:
        # 备用实现
        logloss = float(-np.mean(y * np.log(p_clip) + (1 - y) * np.log(1 - p_clip)))

    # brier score
    brier = float(np.mean((p_clip - y) ** 2))

    # ECE 计算（支持 uniform 或 quantile）
    if bin_strategy == "quantile":
        # quantile bin edges
        try:
            edges = np.unique(np.percentile(p_clip, np.linspace(0, 100, n_bins + 1)))
            if len(edges) - 1 <= 0:
                # fallback to uniform
                bin_idxs = np.minimum((p_clip * n_bins).astype(int), n_bins - 1)
            else:
                # searchsorted to assign bins
                bin_idxs = np.clip(
                    np.searchsorted(edges, p_clip, side="right") - 1, 0, len(edges) - 2
                )
        except Exception:
            bin_idxs = np.minimum((p_clip * n_bins).astype(int), n_bins - 1)
    else:
        bin_idxs = np.minimum((p_clip * n_bins).astype(int), n_bins - 1)

    ece = 0.0
    total = len(y)
    bin_stats = []
    for b in range(n_bins):
        idx = bin_idxs == b
        count = int(idx.sum())
        if count == 0:
            bin_stats.append(
                {"count": 0, "mean_pred": float("nan"), "emp_freq": float("nan")}
            )
            continue
        mean_pred = float(p_clip[idx].mean())
        emp_freq = float(y[idx].mean())
        ece += abs(mean_pred - emp_freq) * count
        bin_stats.append({"count": count, "mean_pred": mean_pred, "emp_freq": emp_freq})
    ece = float(ece / total) if total > 0 else float("nan")

    # accuracy
    acc = float(np.mean((p_clip >= 0.5) == (y == 1)))

    # 校准拟合: 使用 LogisticRegression 拟合 logit(E[y]) = alpha + beta * logit(p)
    X = sp_logit(p_clip).reshape(-1, 1)
    clf = LogisticRegression(C=1e6, solver="lbfgs", max_iter=200)
    clf.fit(X, y)
    alpha = float(clf.intercept_[0])
    beta = float(clf.coef_[0][0])

    return {
        "logloss": logloss,
        "brier": brier,
        "ece": ece,
        "accuracy": acc,
        "reg_alpha": alpha,
        "reg_beta": beta,
        # 'ece_bins': bin_stats,
        "n_samples": int(total),
    }


if __name__ == "__main__":
    df = pd.read_feather("data/p_res.feather")
    df["win_prob"] = df["power_p"]
    res = compute_metrics(df)
    print(res)