TradeEnv/batch_process_trade.py

"""
批量处理 Trade 测试配置文件
自动运行所有测试配置并保存结果
"""

import os
import json
import sys
from datetime import datetime
from typing import Dict, List
import glob

# 添加当前目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

# 直接导入同目录下的模块
from auto_trade_solver import run_auto_trade
import progress_manager


def find_all_trade_configs(test_data_dir: str) -> List[str]:
    """查找所有交易测试配置文件"""
    pattern = os.path.join(test_data_dir, "test_trade_config_*.json")
    config_files = sorted(glob.glob(pattern), key=lambda x: int(os.path.basename(x).split('_')[-1].split('.')[0]))
    return config_files


def process_all_configs(test_data_dir: str, output_dir: str = "output", 
                       user_id: str = "batch_auto", save_dir: str = "user_progress",
                       verbose: bool = False) -> Dict:
    """
    批量处理所有配置文件
    
    Args:
        test_data_dir: 测试数据目录
        output_dir: 结果输出目录（相对于脚本所在目录）
        user_id: 用户ID（用于保存进度）
        save_dir: 进度保存目录
        verbose: 是否打印详细信息
        
    Returns:
        处理结果汇总
    """
    # 获取脚本所在目录
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 构建测试数据目录路径（相对于脚本目录）
    if not os.path.isabs(test_data_dir):
        test_data_dir = os.path.join(script_dir, test_data_dir)
    
    # 构建输出目录路径（相对于脚本目录）
    if not os.path.isabs(output_dir):
        output_dir = os.path.join(script_dir, output_dir)
    
    # 构建保存目录路径（相对于脚本目录）
    if not os.path.isabs(save_dir):
        save_dir = os.path.join(script_dir, save_dir)
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 查找所有配置文件
    config_files = find_all_trade_configs(test_data_dir)
    
    if not config_files:
        print(f"❌ 未找到任何配置文件在 {test_data_dir}")
        return {}
    
    print(f"📊 找到 {len(config_files)} 个配置文件")
    print("=" * 60)
    
    # 用于存储标准格式的环境数据
    environments_data = {}
    results = []
    total_profit_ratio = 0.0
    successful_runs = 0
    failed_runs = 0
    
    for idx, config_path in enumerate(config_files, 1):
        config_name = os.path.basename(config_path)
        env_idx = idx - 1  # 环境索引从0开始
        
        print(f"\n[{idx}/{len(config_files)}] 处理: {config_name}")
        print("-" * 60)
        
        try:
            # 读取配置获取基本信息
            with open(config_path, 'r') as f:
                cfg = json.load(f)
            
            num_stocks = len(cfg.get('stocks', []))
            num_factors = len(cfg.get('variables', []))
            num_days = cfg.get('num_days', 0)
            initial_cash = cfg.get('initial_cash', 0)
            
            print(f"  股票数: {num_stocks}, 因子数: {num_factors}, 天数: {num_days}")
            
            # 运行自动交易（限制为120天）
            result = run_auto_trade(
                config_path,
                verbose=verbose,
                user_id=user_id,
                save_dir=save_dir,
                env_idx=env_idx,
                load_progress=False,  # 批量处理时从头开始
                save_interval=0,  # 批量处理时不自动保存（最后统一保存）
                max_days=120  # 限制为120天
            )
            
            # 从进度文件中读取完整的环境信息
            env_progress = progress_manager.get_task_environment_progress(
                user_id, save_dir, "trade", env_idx
            )
            
            if env_progress:
                # 按照标准格式组织环境数据
                env_key = str(env_idx)
                environments_data[env_key] = {
                    "user_id": env_progress.get("user_id", user_id),
                    "env_idx": env_progress.get("env_idx", env_idx),
                    "env_idx_display": env_progress.get("env_idx_display", env_idx + 1),
                    "day": env_progress.get("day", result.get('total_days', 0)),
                    "cash": env_progress.get("cash", 0),
                    "positions": env_progress.get("positions", []),
                    "prices": env_progress.get("prices", []),
                    "variables_state": env_progress.get("variables_state", []),
                    "history": env_progress.get("history", []),
                    "num_steps": env_progress.get("num_steps", result.get('total_days', 0)),
                    "done": env_progress.get("done", False),
                    "success": env_progress.get("success", False)
                }
            else:
                # 如果无法读取进度文件，使用基本信息
                env_key = str(env_idx)
                environments_data[env_key] = {
                    "user_id": user_id,
                    "env_idx": env_idx,
                    "env_idx_display": env_idx + 1,
                    "day": result.get('total_days', 0),
                    "cash": 0,
                    "positions": [],
                    "prices": [],
                    "variables_state": [],
                    "history": [],
                    "num_steps": result.get('total_days', 0),
                    "done": False,
                    "success": False
                }
            
            # 记录结果（用于汇总报告）
            result_summary = {
                "config_file": config_name,
                "env_idx": env_idx,
                "num_stocks": num_stocks,
                "num_factors": num_factors,
                "num_days": num_days,
                "initial_cash": initial_cash,
                "initial_value": result.get('initial_value', 0),
                "final_value": result.get('final_value', 0),
                "profit": result.get('profit', 0),
                "profit_ratio": result.get('profit_ratio', 0),
                "total_days": result.get('total_days', 0),
                "observation_days": result.get('observation_days', 0),
                "success": True,
                "error": None
            }
            
            results.append(result_summary)
            total_profit_ratio += result.get('profit_ratio', 0)
            successful_runs += 1
            
            print(f"  ✅ 完成: 收益率 {result.get('profit_ratio', 0):.2f}%")
            
        except Exception as e:
            print(f"  ❌ 失败: {str(e)}")
            failed_runs += 1
            
            result_summary = {
                "config_file": config_name,
                "env_idx": env_idx,
                "success": False,
                "error": str(e)
            }
            results.append(result_summary)
    
    # 生成标准格式的输出（参考 lqk.json 格式）
    standard_output = {
        "trade": {
            "environments": environments_data
        }
    }
    
    # 生成汇总报告（保留原有格式用于兼容）
    summary = {
        "timestamp": datetime.now().isoformat(),
        "total_configs": len(config_files),
        "successful_runs": successful_runs,
        "failed_runs": failed_runs,
        "average_profit_ratio": total_profit_ratio / successful_runs if successful_runs > 0 else 0,
        "results": results
    }
    
    # 保存标准格式结果到JSON文件
    output_file = os.path.join(output_dir, f"batch_trade_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(standard_output, f, ensure_ascii=False, indent=2)
    
    # 同时保存汇总报告（可选，用于兼容）
    summary_file = os.path.join(output_dir, f"batch_trade_summary_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
    with open(summary_file, 'w', encoding='utf-8') as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)
    
    # 生成文本报告
    report_file = os.path.join(output_dir, f"batch_trade_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")
    with open(report_file, 'w', encoding='utf-8') as f:
        f.write("=" * 60 + "\n")
        f.write("批量交易测试结果报告\n")
        f.write("=" * 60 + "\n\n")
        f.write(f"生成时间: {summary['timestamp']}\n")
        f.write(f"总配置文件数: {summary['total_configs']}\n")
        f.write(f"成功运行: {summary['successful_runs']}\n")
        f.write(f"失败运行: {summary['failed_runs']}\n")
        f.write(f"平均收益率: {summary['average_profit_ratio']:.2f}%\n\n")
        
        f.write("-" * 60 + "\n")
        f.write("详细结果:\n")
        f.write("-" * 60 + "\n\n")
        
        for result in results:
            f.write(f"配置文件: {result['config_file']}\n")
            if result.get('success'):
                f.write(f"  环境索引: {result.get('env_idx', 'N/A')}\n")
                f.write(f"  股票数: {result.get('num_stocks', 'N/A')}, 因子数: {result.get('num_factors', 'N/A')}\n")
                f.write(f"  初始价值: {result.get('initial_value', 0):.2f}\n")
                f.write(f"  最终价值: {result.get('final_value', 0):.2f}\n")
                f.write(f"  总收益: {result.get('profit', 0):.2f}\n")
                f.write(f"  收益率: {result.get('profit_ratio', 0):.2f}%\n")
                f.write(f"  运行天数: {result.get('total_days', 0)}\n")
            else:
                f.write(f"  ❌ 失败: {result.get('error', 'Unknown error')}\n")
            f.write("\n")
        
        # 收益率排名
        successful_results = [r for r in results if r.get('success')]
        if successful_results:
            f.write("-" * 60 + "\n")
            f.write("收益率排名 (Top 10):\n")
            f.write("-" * 60 + "\n\n")
            sorted_results = sorted(successful_results, key=lambda x: x.get('profit_ratio', 0), reverse=True)
            for i, result in enumerate(sorted_results[:10], 1):
                f.write(f"{i}. {result['config_file']}: {result.get('profit_ratio', 0):.2f}%\n")
    
    print("\n" + "=" * 60)
    print("📊 批量处理完成！")
    print("=" * 60)
    print(f"总配置文件数: {len(config_files)}")
    print(f"成功运行: {successful_runs}")
    print(f"失败运行: {failed_runs}")
    if successful_runs > 0:
        print(f"平均收益率: {total_profit_ratio / successful_runs:.2f}%")
    print(f"\n结果已保存到:")
    print(f"  标准格式 JSON: {output_file}")
    print(f"  汇总报告 JSON: {summary_file}")
    print(f"  文本报告: {report_file}")
    
    return standard_output


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="批量处理 Trade 测试配置文件")
    parser.add_argument("--test_data_dir", type=str, default="../test_data/trade",
                       help="测试数据目录路径")
    parser.add_argument("--output_dir", type=str, default="output",
                       help="结果输出目录")
    parser.add_argument("--user_id", type=str, default="batch_auto",
                       help="用户ID（用于保存进度）")
    parser.add_argument("--save_dir", type=str, default="user_progress",
                       help="进度保存目录")
    parser.add_argument("--verbose", action="store_true",
                       help="打印详细信息")
    
    args = parser.parse_args()
    
    # 处理所有配置
    summary = process_all_configs(
        test_data_dir=args.test_data_dir,
        output_dir=args.output_dir,
        user_id=args.user_id,
        save_dir=args.save_dir,
        verbose=args.verbose
    )