"""
ResearchAgent — 因子发现与评估。

遍历注册因子，计算 IC/IC_IR/分层收益，输出 Top 因子。
"""

import numpy as np
import pandas as pd

from agents.base import BaseAgent


class ResearchAgent(BaseAgent):
    """因子发现 Agent。"""

    name = "Research"
    description = "因子发现与评估"

    def execute(
        self,
        factor_names: list[str] | None = None,
        ts_codes: list[str] | None = None,
        lookahead: int = 5,
        top_n: int = 10,
    ) -> dict:
        """
        遍历因子，计算评估指标。

        参数:
            factor_names: 待评估因子列表，None=全部注册因子
            ts_codes: 股票列表，None=DataManager 全部
            lookahead: 前向收益窗口
            top_n: 返回 Top N 因子

        返回:
            {"top_factors": [...], "all_results": DataFrame, "evaluated": int}
        """
        from factors.registry import list_factors, get_factor

        if factor_names is None:
            # 只评估技术+基本面因子（情绪因子需要额外数据）
            categories = ["动量", "RSI", "MACD", "量价", "布林", "ATR", "均线", "波动率", "换手率", "振幅", "基本面"]
            factor_names = []
            for cat in categories:
                factor_names.extend(list_factors(cat))

        if ts_codes is None:
            stocks = self.dm.get_stock_list()
            # 默认选 50 只代表股（前50只 + 避免API过载）
            ts_codes = list(stocks.index[:50]) if self.dm else ["000001.SZ", "600519.SH"]

        self.log(f"评估 {len(factor_names)} 个因子 × {len(ts_codes)} 只股票")

        results = []
        for fn in factor_names:
            metrics = self._evaluate_factor(fn, ts_codes, lookahead)
            if metrics:
                results.append(metrics)
            self.log(f"  {fn}: IC={metrics.get('ic_mean', 0):.4f}" if metrics else f"  {fn}: SKIP")

        if not results:
            return {"top_factors": [], "all_results": pd.DataFrame(), "evaluated": 0}

        df = pd.DataFrame(results).sort_values("score", ascending=False)
        top = df.head(top_n)

        return {
            "top_factors": top.to_dict("records"),
            "all_results": df,
            "evaluated": len(results),
        }

    def _evaluate_factor(
        self, factor_name: str, ts_codes: list[str], lookahead: int
    ) -> dict | None:
        """对单个因子计算 IC/IC_IR。"""
        from factors.registry import get_factor
        from models.features import FeatureEngine

        try:
            factor = get_factor(factor_name)
        except KeyError:
            return None

        fe = FeatureEngine(lookahead=lookahead, label_type="regression")

        ics = []
        long_rets = []
        short_rets = []
        success = 0

        for ts_code in ts_codes:
            try:
                daily = self.dm.get_daily(ts_code)
                if daily is None or daily.empty:
                    continue
                daily = daily.set_index("trade_date").sort_index()

                factor_df = self.fe.compute(ts_code, [factor])
                if factor_df is None or factor_df.empty:
                    continue

                X, y = fe.build(factor_df, daily, fit=True)
                if X.empty or y.empty or factor_name not in X.columns:
                    continue

                fv = X[factor_name].dropna()
                yv = y.loc[fv.index]
                if len(fv) < 30:
                    continue

                ic = fv.corr(yv, method="spearman")
                ics.append(ic)

                # 分层收益
                top_idx = fv.nlargest(int(len(fv) * 0.2)).index
                bot_idx = fv.nsmallest(int(len(fv) * 0.2)).index
                long_rets.append(yv.loc[top_idx.intersection(yv.index)].mean())
                short_rets.append(yv.loc[bot_idx.intersection(yv.index)].mean())

                success += 1
            except Exception:
                continue

        if success < 3:
            return None

        ic_series = pd.Series(ics)
        ic_mean = ic_series.mean()
        ic_std = ic_series.std()
        icir = ic_mean / ic_std if ic_std > 0 else 0

        # 综合得分 = IC × IC_IR 加权
        score = abs(ic_mean) * max(icir, 0)

        return {
            "name": factor_name,
            "ic_mean": round(float(ic_mean), 4),
            "ic_std": round(float(ic_std), 4),
            "icir": round(float(icir), 4),
            "long_ret": round(float(np.mean(long_rets)), 2) if long_rets else 0,
            "short_ret": round(float(np.mean(short_rets)), 2) if short_rets else 0,
            "stocks_evaluated": success,
            "score": round(float(score), 4),
        }