"""
情绪弹性特征 (4维) — calculate_relaxed_emotion_features(df, market_regime)
180 日窗口，筛选 advance_ratio < 0.20 恐慌日 + advance_ratio > 0.55 贪婪日。
"""
import numpy as np
import pandas as pd
from core.scoring.config import PANIC_ADVANCE_RATIO, GREED_ADVANCE_RATIO, WINDOW_180D


def calculate_relaxed_emotion_features(ctx) -> pd.DataFrame:
    """计算情绪弹性特征 (4维)"""
    kline = ctx.kline.copy()
    mkt = ctx.market_regime.copy() if ctx.market_regime is not None else pd.DataFrame()

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

    # 筛选恐慌日和贪婪日
    panic_dates = set()
    greed_dates = set()

    for _, row in mkt.iterrows():
        td = row['trade_date']
        ar = row.get('advance_ratio', 0.5)
        if ar < PANIC_ADVANCE_RATIO:
            panic_dates.add(td.date() if hasattr(td, 'date') else td)
        if ar > GREED_ADVANCE_RATIO:
            greed_dates.add(td.date() if hasattr(td, 'date') else td)

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

    candidates = ctx.candidates

    # 计算全市场平均振幅 (用于恐慌/贪婪 ratio)
    kline['amp'] = (kline['high'] - kline['low']) / np.where(
        kline['open'] > 0, kline['open'], 1
    )
    market_amp = kline.groupby('td')['amp'].mean().to_dict()

    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(WINDOW_180D)
        closes = g['close'].values
        opens = g['open'].values
        highs = g['high'].values
        lows = g['low'].values
        dates = g['td'].values
        amps = (highs - lows) / np.where(opens > 0, opens, 1)

        feat = {}

        # 恐慌日分析
        panic_indices = [i for i, d in enumerate(dates) if d in panic_dates]
        if panic_indices:
            panic_amp_ratios = []
            panic_drop_ratios = []
            panic_rebounds = []

            for pi in panic_indices:
                td = dates[pi]
                mkt_amp_val = market_amp.get(td, amps[pi])
                if mkt_amp_val > 0:
                    panic_amp_ratios.append(amps[pi] / mkt_amp_val)

                # 恐慌日跌幅
                if pi > 0 and closes[pi - 1] > 0:
                    drop = (closes[pi] - closes[pi - 1]) / closes[pi - 1]
                    # 全市场跌幅: 用 advance_ratio 估算
                    ar = _get_advance_ratio(mkt, td)
                    market_drop = -0.02 if ar < PANIC_ADVANCE_RATIO else -0.005
                    if market_drop != 0:
                        panic_drop_ratios.append(drop / market_drop)

                # 恐慌次日反弹
                if pi + 1 < len(g) and closes[pi] > 0:
                    panic_rebounds.append(highs[pi + 1] / closes[pi] - 1)

            # 64. relaxed_panic_amplitude_ratio
            feat['relaxed_panic_amplitude_ratio'] = (
                np.median(panic_amp_ratios) if panic_amp_ratios else 1.0
            )

            # 65. relaxed_panic_rebound_strength
            feat['relaxed_panic_rebound_strength'] = (
                np.median(panic_rebounds) if panic_rebounds else 0.0
            )
        else:
            feat['relaxed_panic_amplitude_ratio'] = 1.0
            feat['relaxed_panic_rebound_strength'] = 0.0

        # 贪婪日分析
        greed_indices = [i for i, d in enumerate(dates) if d in greed_dates]
        if greed_indices:
            greed_gains = []
            greed_amp_ratios = []

            for gi in greed_indices:
                td = dates[gi]
                if gi > 0 and closes[gi - 1] > 0:
                    gain = (closes[gi] - closes[gi - 1]) / closes[gi - 1]
                    # 全市场收益
                    ar = _get_advance_ratio(mkt, td)
                    market_gain = 0.01 if ar > GREED_ADVANCE_RATIO else 0.003
                    if market_gain > 0:
                        greed_gains.append(gain / market_gain)

                mkt_amp_val = market_amp.get(td, amps[gi])
                if mkt_amp_val > 0:
                    greed_amp_ratios.append(amps[gi] / mkt_amp_val)

            # 66. relaxed_greed_relative_gain
            feat['relaxed_greed_relative_gain'] = (
                np.median(greed_gains) if greed_gains else 0.0
            )

            # 67. relaxed_greed_amplitude_ratio
            feat['relaxed_greed_amplitude_ratio'] = (
                np.median(greed_amp_ratios) if greed_amp_ratios else 1.0
            )
        else:
            feat['relaxed_greed_relative_gain'] = 0.0
            feat['relaxed_greed_amplitude_ratio'] = 1.0

        features[code] = feat

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


def _get_advance_ratio(mkt_df, trade_date):
    """获取指定日期的 advance_ratio"""
    if mkt_df is None or mkt_df.empty:
        return 0.5
    td = trade_date
    match = mkt_df[mkt_df['trade_date'] == td]
    if not match.empty:
        return match.iloc[0].get('advance_ratio', 0.5)
    return 0.5