"""
大盘独立性特征 (3维) — calculate_independence_features(df)
使用 HS300(000300) 作为 benchmark，最近 20 个交易日 OLS 回归。
"""
import numpy as np
import pandas as pd
from core.scoring.config import WINDOW_20D
from core.scoring.features.v3_2_features import _ols_slope


def calculate_independence_features(ctx) -> pd.DataFrame:
    """计算大盘独立性特征 (3维)"""
    kline = ctx.kline.copy()
    hs300 = ctx.hs300_kline.copy() if ctx.hs300_kline is not None else pd.DataFrame()

    if kline.empty or hs300.empty:
        return pd.DataFrame()

    # 准备 HS300 收益率序列
    hs300 = hs300.sort_values('trade_date')
    hs300['market_return'] = hs300['close'].pct_change()
    hs300['market_amplitude'] = (hs300['high'] - hs300['low']) / hs300['open']
    hs300 = hs300.dropna(subset=['market_return', 'market_amplitude'])

    # 对齐日期
    hs300_dates = set(hs300['trade_date'].dt.date
                      if hasattr(hs300['trade_date'], 'dt') else hs300['trade_date'])

    kline = kline.sort_values(['stock_code', 'trade_date'])
    kline['trade_date_dt'] = (kline['trade_date'].dt.date
                              if hasattr(kline['trade_date'], 'dt')
                              else pd.to_datetime(kline['trade_date']).dt.date)

    candidates = ctx.candidates

    # 计算 HS300 最近20日平均振幅
    hs300_tail = hs300.tail(WINDOW_20D)
    hs300_avg_amp = hs300_tail['market_amplitude'].mean() if len(hs300_tail) > 0 else 0

    features = {}
    grouped = kline.groupby('stock_code')

    for code, group in grouped:
        if code not in candidates:
            continue
        if len(group) < 20:
            continue

        g = group.sort_values('trade_date').tail(120)
        closes = g['close'].values

        # 计算个股日收益率
        stock_rets = np.diff(closes) / np.where(closes[:-1] > 0, closes[:-1], 1)

        # 对齐 HS300 收益率 (取对应日期)
        # 简化: 取最近 N 个交易日的数据点
        n = min(WINDOW_20D, len(stock_rets))

        # 获取 HS300 最近 n 天的 market_return
        market_rets = hs300['market_return'].tail(n + 1).values
        if len(market_rets) < n:
            market_rets = hs300['market_return'].values[-n - 1:]

        # 对齐长度
        min_len = min(n, len(market_rets) - 1, len(stock_rets))
        if min_len < 5:  # 至少需要5个数据点做回归
            continue

        stock_ret_window = stock_rets[-min_len:]
        market_ret_window = market_rets[-min_len:]

        feat = {}

        # OLS 回归: stock_ret ~ market_return
        try:
            slope, r_value = _ols_slope(market_ret_window, stock_ret_window)
            residuals = stock_ret_window - slope * market_ret_window

            # 58. market_residual_volatility_20d: std(residuals) × √252
            feat['market_residual_volatility_20d'] = np.std(residuals, ddof=1) * np.sqrt(252)

            # 59. market_independence_ratio_20d: 1 - R²
            feat['market_independence_ratio_20d'] = 1 - r_value ** 2

        except Exception:
            feat['market_residual_volatility_20d'] = 0
            feat['market_independence_ratio_20d'] = 1

        # 60. market_amplitude_deviation_20d: 个股平均振幅 - HS300 平均振幅
        g_amps = (g['high'].values[-20:] - g['low'].values[-20:]) / np.where(
            g['open'].values[-20:] > 0, g['open'].values[-20:], 1
        )
        feat['market_amplitude_deviation_20d'] = np.mean(g_amps) - hs300_avg_amp

        features[code] = feat

    return pd.DataFrame.from_dict(features, orient='index')