#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
通用自动化回测脚本
自动保存结果到：
1. 全局目录: /volume1/stock/sanguo_vnpy/backtest_results/
2. 策略目录: ./strategies/{your_strategy}/backtest_results/ 便于和策略放在一起查看
"""

import sys
import os
import argparse
from datetime import datetime
from pathlib import Path

# 添加路径
sys.path.insert(0, '/volume1/stock/sanguo_vnpy')
sys.path.insert(0, '/volume1/stock/sanguo_vnpy/strategies')

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from vnpy.trader.constant import Interval
from vnpy_ctabacktester import BacktesterEngine


def run_backtest(
    strategy_class,
    strategy_name: str,
    vt_symbol: str = "IF888.CFFEX",
    interval: Interval = Interval.DAILY,
    start_date: str = "20200101",
    end_date: str = "20251231",
    capital: int = 1000000,
    rate: float = 0.3 / 10000,
    slippage: float = 0.2,
    size: int = 300,
    pricetick: float = 0.2,
    strategy_config: dict = None,
):
    """运行单个策略回测"""
    
    # 创建结果保存目录
    # 1. 全局目录
    GLOBAL_RESULT_DIR = Path('/volume1/stock/sanguo_vnpy/backtest_results')
    GLOBAL_RESULT_DIR.mkdir(exist_ok=True, parents=True)
    
    # 2. 策略目录（和策略放在一起便于查看）
    # 找到策略目录
    strategy_file = sys.modules[strategy_class.__module__].__file__
    strategy_dir = Path(strategy_file).parent
    STRATEGY_RESULT_DIR = strategy_dir / "backtest_results"
    STRATEGY_RESULT_DIR.mkdir(exist_ok=True, parents=True)
    
    print("=" * 60)
    print(f"  开始回测: {strategy_name}")
    print(f"  结果保存:")
    print(f"    全局: {GLOBAL_RESULT_DIR}")
    print(f"    策略: {STRATEGY_RESULT_DIR}")
    print("=" * 60)
    print()
    
    # 打印配置
    print(f"📋 回测配置:")
    print(f"  标的: {vt_symbol}")
    print(f"  周期: {interval}")
    print(f"  时间: {start_date} - {end_date}")
    print(f"  初始资金: {capital:,}")
    print()
    
    # 创建回测引擎
    print("🚀 创建回测引擎...")
    engine = BacktesterEngine()
    
    # 添加策略
    print("🧩 添加策略...")
    if strategy_config is None:
        strategy_config = {}
    engine.add_strategy(strategy_class, strategy_config)
    
    # 运行回测
    print("🔄 开始运行回测...")
    start_time = datetime.now()
    
    engine.run_backtest(
        vt_symbol=vt_symbol,
        interval=interval,
        start=start_date,
        end=end_date,
        rate=rate,
        slippage=slippage,
        size=size,
        pricetick=pricetick,
        capital=capital
    )
    
    end_time = datetime.now()
    duration = (end_time - start_time).total_seconds()
    
    print(f"✅ 回测完成！耗时: {duration:.2f} 秒")
    print()
    
    # 获取结果文本
    import io
    from contextlib import redirect_stdout
    
    output = io.StringIO()
    with redirect_stdout(output):
        engine.show_results()
    result_text = output.getvalue()
    
    # 保存结果到全局目录
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    global_result_file = GLOBAL_RESULT_DIR / f"{strategy_name}_{timestamp}.txt"
    
    with open(global_result_file, 'w', encoding='utf-8') as f:
        f.write(f"{strategy_name} 回测结果\n")
        f.write("=" * 50 + "\n")
        f.write(f"时间: {timestamp}\n")
        f.write(f"标的: {vt_symbol}\n")
        f.write(f"周期: {interval}\n")
        f.write(f"时间范围: {start_date} - {end_date}\n")
        f.write(f"初始资金: {capital:,}\n")
        f.write("\n")
        f.write(result_text)
    
    print(f"💾 全局结果已保存: {global_result_file}")
    
    # 保存结果到策略目录
    strategy_result_file = STRATEGY_RESULT_DIR / f"{strategy_name}_{timestamp}.txt"
    with open(strategy_result_file, 'w', encoding='utf-8') as f:
        f.write(f"{strategy_name} 回测结果\n")
        f.write("=" * 50 + "\n")
        f.write(f"时间: {timestamp}\n")
        f.write(f"标的: {vt_symbol}\n")
        f.write(f"周期: {interval}\n")
        f.write(f"时间范围: {start_date} - {end_date}\n")
        f.write(f"初始资金: {capital:,}\n")
        f.write("\n")
        f.write(result_text)
    
    print(f"📂 策略结果已保存: {strategy_result_file}")
    print()
    
    # 保存图表
    print("📈 绘制资金曲线...")
    plt.figure(figsize=(12, 6))
    engine.plot_chart()
    
    # 保存到全局目录
    global_chart_file = GLOBAL_RESULT_DIR / f"{strategy_name}_{timestamp}.png"
    plt.savefig(global_chart_file)
    print(f"📊 全局图表已保存: {global_chart_file}")
    
    # 保存到策略目录
    strategy_chart_file = STRATEGY_RESULT_DIR / f"{strategy_name}_{timestamp}.png"
    plt.savefig(strategy_chart_file)
    print(f"📊 策略图表已保存: {strategy_chart_file}")
    
    print()
    print("=" * 60)
    print(f"  🎉 回测完成！")
    print(f"  📄 全局结果: {global_result_file}")
    print(f"  📂 策略结果: {strategy_result_file}")
    print("=" * 60)
    
    return global_result_file, strategy_result_file


def main():
    """主函数"""
    
    parser = argparse.ArgumentParser(description='自动化运行策略回测')
    parser.add_argument('strategy', help='策略模块名，例如: strategies.guanyu_value_tech_strategy')
    parser.add_argument('--symbol', default='IF888.CFFEX', help='回测标的')
    parser.add_argument('--start', default='20200101', help='开始日期')
    parser.add_argument('--end', default='20251231', help='结束日期')
    parser.add_argument('--capital', type=int, default=1000000, help='初始资金')
    
    args = parser.parse_args()
    
    # 导入策略模块
    import importlib
    module = importlib.import_module(args.strategy)
    
    # 找到策略类（假设第一个类就是策略类）
    strategy_class = None
    for name, obj in module.__dict__.items():
        if isinstance(obj, type) and 'Strategy' in name:
            strategy_class = obj
            break
    
    if strategy_class is None:
        print(f"❌ 在模块 {args.strategy} 中没找到策略类")
        sys.exit(1)
    
    strategy_name = strategy_class.__name__
    
    # 获取策略配置（如果有）
    strategy_config = {}
    if hasattr(module, 'STRATEGY_CONFIG'):
        strategy_config = module.STRATEGY_CONFIG
    
    # 运行回测
    run_backtest(
        strategy_class=strategy_class,
        strategy_name=strategy_name,
        vt_symbol=args.symbol,
        start_date=args.start,
        end_date=args.end,
        capital=args.capital,
        strategy_config=strategy_config,
    )


if __name__ == "__main__":
    main()