# sanguo_factor/metrics.py
"""批量评估指标:全向量化(RankIC=秩相关逐行),不依赖 alphalens.

所有矩阵约定:pandas DataFrame,index=DatetimeIndex(日),columns=vt_symbol,
值=因子值或前瞻收益,NaN=缺失自动从当日截面剔除。
"""
import warnings

import numpy as np
import pandas as pd
from scipy import stats

TRADING_DAYS_PER_YEAR = 244


def _row_pearson(a: np.ndarray, b: np.ndarray, index: pd.Index) -> pd.Series:
    """逐行 Pearson(输入已是秩),全 NaN/无方差行 → NaN."""
    mask = ~(np.isnan(a) | np.isnan(b))
    n = mask.sum(axis=1)
    a0 = np.where(mask, a, np.nan)
    b0 = np.where(mask, b, np.nan)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)
        am = np.nanmean(a0, axis=1, keepdims=True)
        bm = np.nanmean(b0, axis=1, keepdims=True)
    ad = np.where(mask, a0 - am, 0.0)
    bd = np.where(mask, b0 - bm, 0.0)
    denom = np.sqrt((ad ** 2).sum(axis=1) * (bd ** 2).sum(axis=1))
    with np.errstate(invalid="ignore", divide="ignore"):
        ic = np.where(denom > 0, (ad * bd).sum(axis=1) / denom, np.nan)
    ic = np.where(n >= 3, ic, np.nan)  # <3 只无意义
    return pd.Series(ic, index=index)


def rank_corr_rows(A: pd.DataFrame, B: pd.DataFrame) -> pd.Series:
    """逐日 Spearman:先各自按行取秩再逐行 Pearson."""
    cols = A.columns.intersection(B.columns)
    idx = A.index.intersection(B.index)
    a_vals = A.loc[idx, cols].to_numpy(dtype=float)
    b_vals = B.loc[idx, cols].to_numpy(dtype=float)

    # scipy.stats.rankdata ranks NaN positions (NaN sorts last)
    # pandas keeps NaN as NaN, so we need to mask and restore
    a_nan = np.isnan(a_vals)
    b_nan = np.isnan(b_vals)

    # Fill NaN with large negative value (sorts first, then rankdata puts it last)
    # But we need it to sort last, so use large positive
    a_filled = np.where(a_nan, np.inf, a_vals)
    b_filled = np.where(b_nan, np.inf, b_vals)

    # scipy.stats.rankdata with axis=1 ranks per row
    a = stats.rankdata(a_filled, axis=1)
    b = stats.rankdata(b_filled, axis=1)

    # Restore NaN at original NaN positions
    a[a_nan] = np.nan
    b[b_nan] = np.nan

    return _row_pearson(a, b, idx)


def factor_turnover(F: pd.DataFrame) -> float:
    """换手率 = 1 - 相邻两日因子秩相关均值."""
    if len(F) < 2:
        return 0.0
    # Shift indices to align consecutive days
    F_next = F.iloc[1:].reset_index(drop=True)
    F_prev = F.iloc[:-1].reset_index(drop=True)
    corr = rank_corr_rows(F_next, F_prev).replace([np.inf, -np.inf], np.nan).dropna()
    return float(1.0 - corr.mean()) if len(corr) else 0.0


def _quantile_mask(F: pd.DataFrame, top: bool) -> pd.DataFrame:
    """按行把因子值分位选mask(基于升序秩/当日有效数). top=True选最大10%,False选最小10%."""
    vals = F.to_numpy(dtype=float)
    nan_mask = np.isnan(vals)

    # For descending rank (1=max), rank negative values
    # scipy.stats.rankdata sorts ascending by default
    vals_filled = np.where(nan_mask, -np.inf, -vals)  # NaN -> -inf (smallest), negate for descending
    ranks = stats.rankdata(vals_filled, axis=1)

    # Now ranks are in descending order (1=max), but NaN got rank 1, need to shift
    # Actually, let me think more carefully:
    # - pandas rank(axis=1, ascending=False): 1 = max, NaN = NaN
    # - scipy rankdata on negative values: smallest original gets largest rank
    # - So -val (largest original) -> smallest rank value = 1
    # This is what we want!

    # But NaN got filled with -inf, which becomes largest, so rank is 1 - wrong!
    # Need to fill NaN with something that becomes largest rank
    vals_filled = np.where(nan_mask, np.inf, -vals)  # NaN -> +inf (largest after negation)
    ranks = stats.rankdata(vals_filled, axis=1)

    # Now: max original -> -max -> min rank = 1 ✓
    # NaN -> +inf -> max rank (last) ✓

    # Restore NaN at original NaN positions
    ranks[nan_mask] = np.nan

    n = (~nan_mask).sum(axis=1)
    k = np.maximum((n * 0.1).round().astype(int), 1).reshape(-1, 1)

    if top:
        mask = (ranks <= k) & ~nan_mask
    else:
        # For bottom: n - k + 1 to n
        # E.g., n=10, k=1: ranks 10 to 10 (last 1)
        mask = (ranks >= n.reshape(-1, 1) - k + 1) & ~nan_mask

    return pd.DataFrame(mask, index=F.index, columns=F.columns)


def long_short_annual_return(F: pd.DataFrame, R: pd.DataFrame) -> float | None:
    """多空年化:top10% - bottom10% 等权前瞻日收益均值,复利年化."""
    cols = F.columns.intersection(R.columns)
    idx = F.index.intersection(R.index)
    f, r = F.loc[idx, cols], R.loc[idx, cols]
    top = _quantile_mask(f, top=True)
    bot = _quantile_mask(f, top=False)
    daily = (r.where(top).mean(axis=1) - r.where(bot).mean(axis=1)).dropna()
    if daily.empty:
        return None
    return float((1.0 + daily.mean()) ** TRADING_DAYS_PER_YEAR - 1.0)


def decile_annual_returns(F: pd.DataFrame, R: pd.DataFrame) -> list[float | None]:
    """十分组(D1因子最低→D10最高)等权年化收益."""
    cols = F.columns.intersection(R.columns)
    idx = F.index.intersection(R.index)
    f, r = F.loc[idx, cols], R.loc[idx, cols]

    vals = f.to_numpy(dtype=float)
    nan_mask = np.isnan(vals)

    # Fill NaN with large value (sorts last, gets largest rank)
    vals_filled = np.where(nan_mask, np.inf, vals)
    ranks = stats.rankdata(vals_filled, axis=1)

    # Restore NaN at original positions
    ranks[nan_mask] = np.nan

    # Normalize to percentiles
    n = (~nan_mask).sum(axis=1, keepdims=True)
    pct = ranks / n

    out: list[float | None] = []
    for d in range(10):
        sel = (pct > d / 10) & (pct <= (d + 1) / 10)
        daily = r.where(sel).mean(axis=1).dropna()
        out.append(
            float((1.0 + daily.mean()) ** TRADING_DAYS_PER_YEAR - 1.0)
            if len(daily) else None
        )
    return out


def monthly_ic(ic: pd.Series) -> list[dict]:
    """IC 按月聚合(sparkline/详情图数据)."""
    s = ic.dropna()
    if s.empty:
        return []
    g = s.groupby(s.index.to_period("M").to_timestamp()).mean()
    return [{"month": t.strftime("%Y-%m"), "ic": round(float(v), 6)} for t, v in g.items()]


def classify(icir: float | None, t_stat: float | None) -> str:
    """结论信号灯:|ICIR|>=0.3 且 |t|>=2 有效;否则 |t|>=1.5 观察;其余淘汰."""
    if icir is None or t_stat is None:
        return "eliminated"
    if abs(icir) >= 0.3 and abs(t_stat) >= 2.0:
        return "effective"
    if abs(t_stat) >= 1.5:
        return "watch"
    return "eliminated"


def _period_stats(ic: pd.Series) -> dict:
    s = ic.dropna()
    n = len(s)
    if n < 2:
        return {"count": int(n), "ic_mean": None, "ic_std": None,
                "icir": None, "t_stat": None, "win_rate": None,
                "monthly_ic": monthly_ic(ic)}
    m = float(s.mean())
    sd = float(s.std())
    icir = m / sd if sd > 0 else None
    t_stat = m / (sd / n ** 0.5) if sd > 0 else None
    return {
        "count": int(n),
        "ic_mean": m,
        "ic_std": sd,
        "icir": icir,
        "t_stat": t_stat,
        "win_rate": float((s > 0).mean()),
        "monthly_ic": monthly_ic(ic),
    }


def summarize_factor(F: pd.DataFrame, R1: pd.DataFrame, R5: pd.DataFrame, R10: pd.DataFrame) -> dict:
    """单因子全指标:三周期 IC 族 + 换手 + 多空 + 十分组 + 结论."""
    out: dict = {"turnover": factor_turnover(F)}
    for p, R in (("1", R1), ("5", R5), ("10", R10)):
        ic = rank_corr_rows(F, R)
        stats = _period_stats(ic)
        stats["ls_annual"] = long_short_annual_return(F, R)
        stats["deciles"] = decile_annual_returns(F, R)
        stats["conclusion"] = classify(stats["icir"], stats["t_stat"])
        out[p] = stats
    return out