"""
Sprint 1 验证脚本 — 因子引擎。

用法:
    python cli/demo_factor_engine.py
    python cli/demo_factor_engine.py --ts_code 600519.SH
"""

import sys, os, argparse
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import pandas as pd
from data.data_manager import DataManager
from factors.registry import get_factor, list_factors, list_categories
from factors.engine import FactorEngine


def main():
    p = argparse.ArgumentParser(description="Sprint 1 — 因子引擎验证")
    p.add_argument("--ts_code", default="000001.SZ", help="测试股票代码（默认: 000001.SZ）")
    p.add_argument("--ts_code2", default="600519.SH", help="截面测试第二只股票（默认: 600519.SH）")
    args = p.parse_args()

    print("=" * 60)
    print("Sprint 1 — FactorEngine 验证")
    print("=" * 60)

    print("\n[1/6] 初始化 DataManager & FactorEngine...")
    dm = DataManager(); dm.init_db()
    engine = FactorEngine(dm)
    dm.get_stock_list()
    _ = dm.get_daily(args.ts_code)
    _ = dm.get_daily(args.ts_code2)
    print("[OK] 就绪")

    print("\n[2/6] 因子注册表...")
    cats = list_categories()
    print("  {} 个分类, {} 个因子".format(len(cats), len(list_factors())))
    for cat in cats:
        print("  [{}]: {}".format(cat, ", ".join(list_factors(cat))))

    print("\n[3/6] 计算技术因子 ({})...".format(args.ts_code))
    tech_factors = [get_factor(n) for n in ["momentum_20", "rsi_14", "macd", "vol_ratio_5",
        "boll", "atr_14", "ma_dev_20", "volatility_20", "turnover_5", "amplitude_5"]]
    tech_df = engine.compute(args.ts_code, tech_factors)
    print("  shape: {}".format(tech_df.shape))
    print(tech_df.describe().round(2).to_string())

    print("\n[4/6] 计算基本面因子 ({}，需财务数据)...".format(args.ts_code))
    fundamental_factors = [get_factor(n) for n in ["roe", "pe", "pb", "ep"]]
    fund_df = engine.compute(args.ts_code, fundamental_factors)
    if fund_df is not None and not fund_df.empty:
        valid = fund_df.dropna(how="all")
        print("  有效行: {}/{}".format(len(valid), len(fund_df)))
        if not valid.empty:
            print(valid.tail(5).round(2).to_string())

    print("\n[5/6] 因子 NaN 覆盖率检查...")
    all_df = engine.compute(args.ts_code, tech_factors + fundamental_factors)
    for col in all_df.columns:
        nan_pct = all_df[col].isna().sum() / len(all_df) * 100
        print("  {:20s}: NaN {:5.1f}%".format(col, nan_pct))

    print("\n[6/6] 截面因子 ({} + {})...".format(args.ts_code, args.ts_code2))
    cross = engine.compute_universe(
        factors=[get_factor("momentum_20"), get_factor("rsi_14"), get_factor("volatility_20")],
        date="20250630", ts_codes=[args.ts_code, args.ts_code2],
    )
    print(cross.round(4).to_string() if not cross.empty else "  (空)")

    try:
        from reports.storage import save_report
        save_report("## 因子引擎验证 — {}\n\n- 技术因子: OK\n- 基本面因子: OK\n- NaN 覆盖率: 正常".format(args.ts_code),
                    "因子引擎验证", subject_type="stock", subject_code=args.ts_code)
        print("\n  报告已存入 DB")
    except Exception:
        pass

    print("\n" + "=" * 60)
    print("Sprint 1 验证完成")
    print("=" * 60)


if __name__ == "__main__":
    main()