# sanguo_factor/batch_eval.py
"""批量评估引擎:全A bars → 进程内逐因子 calculate_by_expression → 指标 → 落盘.

不走 AlphaDataset.prepare_data(其 spawn 池对每个表达式 pickle 整个 DataFrame,
9M 行 × 258 因子的传输开销不可接受);calculate_by_expression 纯进程内 polars,
内存随单因子天然有界。现有单因子分析链路(analyzer/alphalens tears)零改动.
"""
import sys
import os
import time

_VNPY_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "vnpy_v4.4.0"))
if _VNPY_SRC not in sys.path:
    sys.path.insert(0, _VNPY_SRC)

import pandas as pd
import polars as pl

from .universe import load_universe_bars, evaluation_filter
from .registry import get_factor
from . import eval_store
from .metrics import summarize_factor


def _forward_return_matrices(close_wide: pd.DataFrame, periods=(1, 5, 10)) -> dict[int, pd.DataFrame]:
    out = {}
    for p in periods:
        out[p] = close_wide.shift(-p) / close_wide - 1.0
    return out


def run_batch_eval(
    factor_names: list[str],
    start: str,
    end: str,
    eval_db: str,
    label: str,
    universe: str = "all_a",
    symbols: list[str] | None = None,
    limit: int | None = None,
    cfg=None,
    progress_cb=None,
    vnpy_db_override: str | None = None,
) -> dict:
    """跑一轮批量评估,结果增量写入 eval_db,返回摘要."""
    from vnpy.alpha.dataset.utility import calculate_by_expression

    if cfg is None:
        from sanguo_data.config import load_config, find_config_path
        cfg = load_config(find_config_path())
    vnpy_db = vnpy_db_override or cfg.data_paths["vnpy_db"]

    t0 = time.time()
    bars = load_universe_bars(vnpy_db, start, end, symbols=symbols, limit=limit)
    if bars.height == 0:
        raise ValueError(f"股票池为空: vnpy_db={vnpy_db} window={start}~{end}")

    alpha_df = bars.select(["vt_symbol", "datetime", "open", "high", "low", "close",
                            "volume", "turnover", "open_interest", "vwap"])
    eval_rows = evaluation_filter(bars, start, end).select(["vt_symbol", "datetime"])

    close_wide = (
        bars.select(["datetime", "vt_symbol", "close"])
        .pivot(index="datetime", on="vt_symbol", values="close")
        .sort("datetime").to_pandas().set_index("datetime")
    )
    close_wide.index = pd.to_datetime(close_wide.index)
    rets = _forward_return_matrices(close_wide)

    universe_label = universe if symbols is None else "custom"
    eval_store.init_db(eval_db)
    run_id = eval_store.create_run(
        eval_db, label=label, universe=universe_label,
        symbols_count=bars["vt_symbol"].n_unique(), factors_total=len(factor_names),
        start=start, end=end, params={"limit": limit, "symbols": symbols[:20] if symbols else None},  # 截断防 4400 只全量塞 JSON
    )

    errors: list[str] = []
    done = 0
    buffer: list[dict] = []
    for i, name in enumerate(factor_names):
        row = _eval_one(name, alpha_df, eval_rows, rets, calculate_by_expression)
        if "error" in row["metrics"]:
            errors.append(name)
        buffer.append(row)
        done += 1
        if len(buffer) >= 20 or done == len(factor_names):
            eval_store.save_results(eval_db, run_id, buffer)
            buffer = []
        if progress_cb:
            progress_cb(done, len(factor_names), name)

    eval_store.finish_run(eval_db, run_id, "done", factors_done=done)
    return {
        "run_id": run_id,
        "factors_total": len(factor_names),
        "factors_done": done,
        "errors": errors,
        "elapsed_sec": round(time.time() - t0, 1),
        "symbols_count": int(bars["vt_symbol"].n_unique()),
    }


def _eval_one(name: str, alpha_df: pl.DataFrame, eval_rows: pl.DataFrame,
              rets: dict[int, pd.DataFrame], calculate_by_expression) -> dict:
    factor = get_factor(name)
    if factor is None:
        return {"factor": name, "category": "unknown", "expression": "",
                "metrics": {"error": f"因子未注册: {name}"}}
    try:
        res = calculate_by_expression(alpha_df, factor["expression"])
        f_long = res.join(eval_rows, on=["vt_symbol", "datetime"], how="inner")
        F = (
            f_long.pivot(index="datetime", on="vt_symbol", values="data")
            .sort("datetime").to_pandas().set_index("datetime")
        )
        F.index = pd.to_datetime(F.index)
        if F.empty or len(F.index) == 0:
            return {"factor": name, "category": factor["category"],
                    "expression": factor["expression"],
                    "metrics": {"error": "因子矩阵为空(join 后无评估窗行,检查 datetime dtype/窗口)"}}
        metrics = summarize_factor(F, rets[1], rets[5], rets[10])
        return {"factor": name, "category": factor["category"],
                "expression": factor["expression"], "metrics": metrics}
    except Exception as e:  # 单因子失败不拖垮整批
        return {"factor": name, "category": factor["category"],
                "expression": factor["expression"],
                "metrics": {"error": f"{type(e).__name__}: {e}"}}