"""回测相对/绝对指标计算（empyrical，聚宽同源口径）。纯函数。"""
from dataclasses import dataclass, field
from typing import Dict, Literal
import math
import numpy as np
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 + 基准日收益计算聚宽级指标。

    daily_df: vnpy calculate_result() 产出，须含 "return" 列（日收益率），index 为日期。
    benchmark_returns: 基准日收益率 Series，index 对齐 daily_df。
    period: 年化周期（整数，默认252交易日），empyrical 内部使用 'daily'
    """
    strat = daily_df["return"].astype(float)
    # 对齐
    aligned = pd.concat([strat.rename("s"), benchmark_returns.rename("b")], axis=1).dropna()
    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 保证)
    roll_beta = pd.Series(index=s.index, dtype=float)
    roll_alpha = pd.Series(index=s.index, dtype=float)
    for i in range(len(s)):
        sub = aligned.iloc[: i + 1]
        if len(sub) >= 2 and sub["b"].var() > 0:
            beta = sub["s"].cov(sub["b"]) / sub["b"].var()
            alpha = sub["s"].mean() - beta * sub["b"].mean()
            roll_beta.iloc[i] = beta
            roll_alpha.iloc[i] = alpha * period
    # drawdown: 从峰值回落 (值 <= 0)
    cummax = equity.cummax()
    drawdown = (equity - cummax) / cummax

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