etf/scripts/export_rotation_data.py

"""
导出轮动策略回测所用的原始数据到本地文件夹

导出内容：
1. index_data.csv   - 指数价格数据（宽格式，用于因子计算）
2. etf_data.csv     - ETF价格数据（宽格式，用于收益计算）
3. etf_nav_data.csv - ETF净值数据（宽格式，用于溢价率计算）
4. benchmark_data.csv - 基准数据
5. config_snapshot.yaml - 当时使用的策略配置快照
"""

import sys
import os
import time
import shutil
from datetime import datetime
from pathlib import Path

# 添加项目根目录
sys.path.insert(0, str(Path(__file__).parent.parent))

import yaml
import pandas as pd
from dotenv import load_dotenv

load_dotenv()

from strategies.rotation.engine import RotationStrategy
from config.settings import DEFAULT_BENCHMARK_CODE


def main():
    # 加载配置
    config_path = Path(__file__).parent.parent / "config" / "strategies" / "rotation.yaml"
    with open(config_path, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)

    # 如果未设置 end_date，默认使用今天
    if not config.get("end_date"):
        config["end_date"] = datetime.now().strftime("%Y-%m-%d")

    start_date = config["start_date"]
    end_date = config["end_date"]

    print("=" * 60)
    print("  导出轮动策略回测数据")
    print("=" * 60)
    print(f"  回测区间: {start_date} ~ {end_date}")
    print(f"  候选标的: {len(config.get('code_list', {}))} 只")

    # 创建输出目录
    export_dir = Path(__file__).parent.parent / "data" / "rotation_backtest_data"
    export_dir.mkdir(parents=True, exist_ok=True)
    print(f"  输出目录: {export_dir}")

    # 创建策略实例（仅用于获取数据）
    strategy = RotationStrategy(config)

    # 获取数据
    print("\n" + "=" * 60)
    print("开始下载数据...")
    print("=" * 60)

    benchmark_code = config.get("benchmark", {}).get("code", DEFAULT_BENCHMARK_CODE)
    code_config = config.get("code_list", {})

    with strategy.data_source:
        index_data, etf_data, etf_nav_data, benchmark_data, valid_codes = (
            strategy.data_source.fetch_all(
                code_config, benchmark_code, start_date, end_date
            )
        )

    # 保存数据
    print("\n" + "=" * 60)
    print("保存数据到本地...")
    print("=" * 60)

    saved_files = []

    # 1. 指数价格数据
    if index_data is not None:
        path = export_dir / "index_data.csv"
        index_data.to_csv(path)
        saved_files.append(("index_data.csv", index_data.shape, "指数价格（因子计算用）"))
        print(f"  ✓ index_data.csv: {index_data.shape[0]} 行 × {index_data.shape[1]} 列")

    # 2. ETF价格数据
    if etf_data is not None:
        path = export_dir / "etf_data.csv"
        etf_data.to_csv(path)
        saved_files.append(("etf_data.csv", etf_data.shape, "ETF价格（收益计算用）"))
        print(f"  ✓ etf_data.csv: {etf_data.shape[0]} 行 × {etf_data.shape[1]} 列")

    # 3. ETF净值数据
    if etf_nav_data is not None:
        path = export_dir / "etf_nav_data.csv"
        etf_nav_data.to_csv(path)
        saved_files.append(("etf_nav_data.csv", etf_nav_data.shape, "ETF净值（溢价率计算用）"))
        print(f"  ✓ etf_nav_data.csv: {etf_nav_data.shape[0]} 行 × {etf_nav_data.shape[1]} 列")

    # 4. 基准数据
    if benchmark_data is not None:
        path = export_dir / "benchmark_data.csv"
        benchmark_data.to_csv(path)
        saved_files.append(("benchmark_data.csv", benchmark_data.shape, "基准指数"))
        print(f"  ✓ benchmark_data.csv: {benchmark_data.shape[0]} 行")

    # 5. 有效代码列表
    codes_path = export_dir / "valid_codes.txt"
    with open(codes_path, "w") as f:
        for code in valid_codes:
            name = code_config.get(code, {}).get("name", code)
            etf = code_config.get(code, {}).get("etf", "")
            market = code_config.get(code, {}).get("market", "")
            f.write(f"{code}\t{name}\t{etf or '-'}\t{market}\n")
    print(f"  ✓ valid_codes.txt: {len(valid_codes)} 只有效标的")

    # 6. 策略配置快照
    config_snapshot_path = export_dir / "config_snapshot.yaml"
    shutil.copy2(config_path, config_snapshot_path)
    print(f"  ✓ config_snapshot.yaml: 策略配置快照")

    # 汇总
    print("\n" + "=" * 60)
    print("导出完成！")
    print("=" * 60)
    print(f"  目录: {export_dir}")
    print(f"  文件数: {len(saved_files) + 2}")
    print(f"  数据区间: {start_date} ~ {end_date}")
    print(f"  有效标的: {len(valid_codes)} 只")
    for fname, shape, desc in saved_files:
        print(f"    - {fname}: {shape} ({desc})")


if __name__ == "__main__":
    main()