"""价量因子库。

所有因子只用 T 日及之前的数据计算（严格避免前向偏差），
因子值统一设计为"越大越好"，截面 Z-score 标准化后加权合成。

当前为纯价量因子（日线 OHLCV 可得）。基本面因子（PE/PB/ROE）
需要历史基本面面板数据支持，属于扩展点，见 README。
"""
from __future__ import annotations

from typing import Dict

import numpy as np
import pandas as pd

# 因子注册表：名称 -> 计算函数（输入单只股票 OHLCV，输出按日期索引的因子序列）
FACTOR_FUNCS = {}


def factor(name, lookback=None):
    def deco(fn):
        FACTOR_FUNCS[name] = (fn, lookback)
        return fn
    return deco


@factor("momentum")
def f_momentum(df: pd.DataFrame, lb: int = 60) -> pd.Series:
    """动量：过去 lb 日累计收益（剔除近 5 日，避免短期反转污染，Jegadeesh-Titman 式）。"""
    close = df["close"]
    return close.shift(5) / close.shift(lb) - 1.0


@factor("quality_momentum")
def f_quality_momentum(df: pd.DataFrame, lb: int = 120) -> pd.Series:
    """风险调整动量：区间收益 / 区间波动（夏普式动量，更平滑稳健）。"""
    ret = df["close"].pct_change()
    win = ret.rolling(lb)
    mu, sd = win.mean(), win.std()
    return (mu / sd.replace(0, np.nan)) * np.sqrt(252)


@factor("reversal")
def f_reversal(df: pd.DataFrame, lb: int = 5) -> pd.Series:
    """短期反转：近 lb 日收益取负（跌多了反弹）。"""
    close = df["close"]
    return -(close / close.shift(lb) - 1.0)


@factor("low_vol")
def f_low_vol(df: pd.DataFrame, lb: int = 60) -> pd.Series:
    """低波动：近 lb 日收益波动取负（低波动异象）。"""
    ret = df["close"].pct_change()
    return -ret.rolling(lb).std()


@factor("trend")
def f_trend(df: pd.DataFrame, lb: int = 60) -> pd.Series:
    """趋势强度：价格相对 lb 日均线的偏离。"""
    close = df["close"]
    return close / close.rolling(lb).mean() - 1.0


@factor("high_52w")
def f_high_52w(df: pd.DataFrame, lb: int = 250) -> pd.Series:
    """52 周高点邻近度（George-Hwang 动量）：越接近前高，动量延续概率越高。"""
    close = df["close"]
    return close / close.rolling(lb, min_periods=120).max()


@factor("liquidity")
def f_liquidity(df: pd.DataFrame, lb: int = 20) -> pd.Series:
    """流动性：近 lb 日均成交额（对数）。保证可交易性，避免选中僵尸股。"""
    amount = df["close"] * df["volume"]
    return np.log(amount.rolling(lb).mean().clip(lower=1.0))


def compute_factor_panel(
    data_map: Dict[str, pd.DataFrame],
    factor_name: str,
    lookback: int,
    dates: pd.DatetimeIndex,
) -> pd.DataFrame:
    """对全部标的计算某因子的宽表（index=日期, columns=代码）。"""
    fn, _ = FACTOR_FUNCS[factor_name]
    out = {}
    for code, df in data_map.items():
        try:
            s = fn(df, lookback)
            out[code] = s.reindex(dates)
        except Exception:
            continue
    panel = pd.DataFrame(out)
    return panel


def winsorize_zscore(row: pd.Series, n_mad: float = 3.0) -> float:
    """截面去极值 + Z-score（逐日调用太慢，改用向量化版本，见 cross_section_z）。"""
    raise NotImplementedError


def cross_section_z(panel: pd.DataFrame) -> pd.DataFrame:
    """逐日截面：去极值（MAD 3 倍）后 Z-score 标准化。"""
    med = panel.median(axis=1)
    mad = (panel.sub(med, axis=0)).abs().median(axis=1)
    upper = med + 3.0 * 1.4826 * mad
    lower = med - 3.0 * 1.4826 * mad
    clipped = panel.clip(lower=lower, upper=upper, axis=0)
    mean = clipped.mean(axis=1)
    std = clipped.std(axis=1).replace(0, np.nan)
    return clipped.sub(mean, axis=0).div(std, axis=0)