"""回测相对/绝对指标计算（empyrical，聚宽同源口径）。纯函数。"""
from dataclasses import dataclass, field
from typing import Dict, Literal
import math
import numpy as np
# empyrical 0.5.5 引用了 NumPy 2.0 已移除的别名（np.NINF / np.NaN / np.Inf / np.PINF），
# 不补回会在 sortino_ratio/downside_risk 等函数里抛 AttributeError，导致整块相对指标计算
# 失败、_metrics.json 不生成、结果页回退到 vnpy 原始字段（单位混乱）。
for _alias, _val in (("NINF", -np.inf), ("Inf", np.inf), ("PINF", np.inf),
                     ("NaN", np.nan), ("NAN", np.nan), ("infty", np.inf)):
    if not hasattr(np, _alias):
        setattr(np, _alias, _val)
import pandas as pd
import empyrical

BenchmarkCode = Literal["hs300", "zz500"]
BENCHMARK_SYMBOL: Dict[str, str] = {"hs300": "sh000300", "zz500": "sz000905"}


@dataclass
class MetricsResult:
    scalars: Dict[str, float] = field(default_factory=dict)
    series: Dict[str, pd.Series] = field(default_factory=dict)


def compute_metrics(
    daily_df: pd.DataFrame,
    benchmark_returns: pd.Series,
    period: int = 252,
) -> MetricsResult:
    """对 vnpy daily_df + 基准日收益计算聚宽级指标。

    H4: 从 daily_df["balance"] 自算 simple return（vnpy df["return"] 是 log return，
    empyrical 期望 simple return，直接用会失真）。
    H5: period='daily' → empyrical 内部 252 交易日年化，口径统一。
    benchmark 对齐：reindex 到策略交易日 + ffill，不丢策略日期（原 dropna 会丢停牌日）。

    daily_df: vnpy calculate_result() 产出，须含 "balance" 列（账户余额），index 为日期。
    benchmark_returns: 基准日收益率 Series，index 为日期（不必与 daily_df 对齐）。
    period: 年化周期（整数，默认252交易日），empyrical 内部使用 'daily'
    """
    # H4: simple return from balance（pct_change 首项 NaN → 0）
    if "balance" not in daily_df.columns:
        raise ValueError("daily_df 缺少 balance 列，无法计算 simple return")
    strat = daily_df["balance"].astype(float).pct_change().fillna(0)

    # benchmark ffill 对齐：reindex 到策略交易日，前向填充，不丢策略日期
    b = benchmark_returns.reindex(daily_df.index).ffill().fillna(0)

    aligned = pd.concat([strat.rename("s"), b.rename("b")], axis=1)
    s, b = aligned["s"], aligned["b"]

    scalars = {
        "total_return": float(empyrical.cum_returns_final(s)),
        "annual_return": float(empyrical.annual_return(s, period='daily')),
        "alpha": float(empyrical.alpha(s, b, period='daily')),
        "beta": float(empyrical.beta(s, b)),
        "sharpe_ratio": float(empyrical.sharpe_ratio(s, period='daily')),
        "sortino_ratio": float(empyrical.sortino_ratio(s, period='daily')),
        "information_ratio": float(empyrical.excess_sharpe(s, b)),
        "annual_volatility": float(empyrical.annual_volatility(s, period='daily')),
        "max_drawdown": float(empyrical.max_drawdown(s)),
        "benchmark_return": float(empyrical.cum_returns_final(b)),
        "benchmark_volatility": float(empyrical.annual_volatility(b, period='daily')),
    }
    # Sanitize non-finite floats (NaN/Inf from degenerate inputs) → None for JSON safety
    scalars = {k: (None if isinstance(v, float) and not math.isfinite(v) else v) for k, v in scalars.items()}

    equity = empyrical.cum_returns(s)
    bench_curve = empyrical.cum_returns(b)
    # rolling alpha/beta (expanding window 用于画图，口径由 scalars 保证)
    # 向量化 expanding (O(n)) 替代原 for 循环 (O(n²))——分钟级 daily_df(万行)原循环卡死
    # (cta_c287f4b5 卡 5min CPU88% 根因)。expanding.cov/.var/.mean 数值与原循环一致
    # (ddof=1)；var_b<=0(含单点 NaN)处 beta=NaN，与原 len>=2 且 var>0 守卫等价。
    cov_sb = aligned["s"].expanding().cov(aligned["b"])
    var_b = aligned["b"].expanding().var()
    roll_beta = cov_sb / var_b.where(var_b > 0)
    roll_alpha = (aligned["s"].expanding().mean() - roll_beta * aligned["b"].expanding().mean()) * period
    # drawdown: 从峰值回落 (值 <= 0)
    cummax = equity.cummax()
    drawdown = (equity - cummax) / cummax

    # Rolling annualized volatility (quarterly window) for the volatility chart
    _vol_window = min(63, len(s))
    if _vol_window >= 2:
        vol_strategy = s.rolling(_vol_window, min_periods=2).std() * np.sqrt(period)
        vol_benchmark = b.rolling(_vol_window, min_periods=2).std() * np.sqrt(period)
    else:
        vol_strategy = pd.Series([np.nan] * len(s), index=s.index)
        vol_benchmark = pd.Series([np.nan] * len(s), index=s.index)

    series = {
        "equity_curve": equity,
        "benchmark_curve": bench_curve,
        "alpha": roll_alpha,
        "beta": roll_beta,
        "drawdown": drawdown,
        "volatility_strategy": vol_strategy,
        "volatility_benchmark": vol_benchmark,
    }
    return MetricsResult(scalars=scalars, series=series)