"""
backtest.py - 通用策略回测工具
支持 4 个策略:
  - rsi2_revert: RSI(2) < 10 做多, RSI(2) > 90 做空, MA50 趋势过滤
  - vwap_revert: 价格偏离 VWAP > 1.5σ 回归
  - early_bird: 开盘 30 min 涨跌幅 + 量 > 1.5× → 顺势
  - turtle_breakout: 20 周期突破 + 10 周期反向出场
  - sma_breakout: SMA5 > SMA10 + 价格突破前高 (现有默认)

用法:
  python3 backtest.py --strategy rsi2_revert --symbol NVDA
  python3 backtest.py --strategy turtle_breakout --symbol 0700.HK --days 60

默认 K 线 = 1h, Yahoo Finance 数据源 (美股 NVDA/AAPL 等, 港股 0700.HK 等)
"""

import argparse
import json
import os
import sys
from datetime import datetime, timedelta
from typing import List, Dict, Optional, Tuple

# 本地依赖
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from strategy_registry import get_strategy, list_strategies
from indicators import sma, ema, rsi, atr, vwap, vwap_std, donchian_breakout


# ============== K 线获取 ==============

def fetch_klines_longbridge(symbol: str, days: int = 30, interval: str = '1h') -> Optional[Dict]:
    """LongPort CLI 数据源 (通过 proxychains4 + Clash 香港出口).
    支持美股/港股. Yahoo Finance 国内 VPS 经常 rate limit, 长桥更稳.
    """
    # 港股 ticker Yahoo 是 4 位数字带前导 0, 长桥是 0700.HK
    # 长桥 K 线是 period 内全部, 倒序. 我们重新排序为时间正序.
    bar_map = {
        '1m': '1m', '5m': '5m', '15m': '15m', '30m': '30m',
        '60m': '60m', '1h': '60m', 'day': 'day', '1d': 'day', 'week': 'week',
    }
    bar = bar_map.get(interval, '60m')

    env = os.environ.copy()
    env['LONGBRIDGE_HTTP_URL'] = 'https://openapi.longbridge.com'
    env['LONGBRIDGE_REGION'] = 'ap'
    env['LONGBRIDGE_TRADE_ENABLED'] = 'true'
    cmd = [
        'proxychains4', '-f', os.path.expanduser('~/.proxychains/proxychains.conf'),
        '/home/openclaw/.local/bin/longbridge', '--profile', 'lb_real',
        'candlesticks', symbol, bar, '--json',
    ]
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=30, env=env)
        if r.returncode != 0:
            return None
        import re
        # 找 JSON 起止位置
        json_match = re.search(r'\[\s*\{', r.stdout)
        if not json_match:
            return None
        json_text = '[' + r.stdout[json_match.start()+1:]
        data = json.loads(json_text)
        if not data:
            return None
        # 倒序 → 正序
        data = list(reversed(data))
        return {
            'opens': [k['open'] for k in data],
            'highs': [k['high'] for k in data],
            'lows':  [k['low']  for k in data],
            'closes':[k['close'] for k in data],
            'volumes':[k.get('volume', 0) or 0 for k in data],
            'timestamps':[k.get('timestamp', '') for k in data],
        }
    except Exception as e:
        print(f"⚠️ {symbol} 长桥 K线拉取失败: {e}")
        return None


import subprocess  # 在 fetch_klines_longbridge 后 import

def fetch_klines_yahoo(symbol: str, days: int = 30, interval: str = '1h') -> Optional[Dict]:
    """Yahoo Finance fallback (国内 VPS 可能 rate limit)."""
    try:
        import yfinance as yf
        yahoo_sym = symbol.replace('.US', '').replace('.HK', '.HK')
        df = yf.download(tickers=yahoo_sym, period=f'{days}d',
                          interval=interval, progress=False, auto_adjust=True)
        if df is None or len(df) < 10:
            return None
        if hasattr(df.columns, 'names') and len(df.columns.names) > 1:
            df.columns = df.columns.droplevel(0)
        expected = ['Open', 'High', 'Low', 'Close', 'Volume']
        if not all(col in df.columns for col in expected):
            return None
        return {
            'opens': df['Open'].tolist(),
            'highs': df['High'].tolist(),
            'lows': df['Low'].tolist(),
            'closes': df['Close'].tolist(),
            'volumes': df['Volume'].fillna(0).tolist(),
            'timestamps': df.index.tolist(),
        }
    except Exception as e:
        print(f"⚠️ yahoo {symbol} fallback 失败: {e}")
        return None


def fetch_klines(symbol: str, days: int = 30, interval: str = '1h') -> Optional[Dict]:
    """统一入口: 长桥 → Yahoo fallback."""
    klines = fetch_klines_longbridge(symbol, days, interval)
    if klines:
        return klines
    print("⚠️ 长桥拉数据失败, fallback Yahoo...")
    return fetch_klines_yahoo(symbol, days, interval)


# ============== 策略信号生成 ==============

def signal_rsi2_revert(klines: Dict, params) -> List[Dict]:
    """RSI(2) 超卖反弹信号"""
    closes = klines['closes']
    opens = klines['opens']
    atr_vals = atr(klines['highs'], klines['lows'], closes, 14)
    rsi2 = rsi(closes, params.rsi_period)
    ma50 = sma(closes, params.rsi2_ma_filter)

    signals = []
    cooldown = 0
    # 跳过前面 (50 = MA50 + ATR14 + RSI2 都需要预热)
    start = max(50, params.rsi2_ma_filter + 1)
    for i in range(start, len(closes)):
        cooldown -= 1
        if cooldown > 0:
            continue
        if rsi2[i] is None or ma50[i] is None or atr_vals[i] is None:
            continue

        # 入场
        side = None
        if rsi2[i] < params.rsi_buy_threshold and closes[i] > ma50[i] and opens[i] > closes[i-1]:
            side = 'long'
        elif rsi2[i] > params.rsi_sell_threshold and closes[i] < ma50[i] and opens[i] < closes[i-1]:
            side = 'short'

        if not side:
            continue

        entry = closes[i]
        sl_price = entry - atr_vals[i] * params.sl_atr_multi if side == 'long' else entry + atr_vals[i] * params.sl_atr_multi
        tp_price = entry + atr_vals[i] * params.tp_atr_multi if side == 'long' else entry - atr_vals[i] * params.tp_atr_multi

        signals.append({
            'i': i, 'ts': klines['timestamps'][i], 'side': side,
            'entry': entry, 'sl': sl_price, 'tp': tp_price,
            'atr': atr_vals[i],
        })
        cooldown = params.cooldown_bars
    return signals


def signal_sma_breakout(klines: Dict, params) -> List[Dict]:
    """SMA 突破 (现有默认, 用来对比)"""
    closes = klines['closes']
    opens = klines['opens']
    highs = klines['highs']
    lows = klines['lows']
    atr_vals = atr(highs, lows, closes, 14)

    sma5 = sma(closes, 5)
    sma10 = sma(closes, 10)

    signals = []
    cooldown = 0
    for i in range(15, len(closes)):
        cooldown -= 1
        if cooldown > 0:
            continue
        # SMA 突破: SMA5 > SMA10 + 突破前高
        if sma5[i] is None or sma10[i] is None or atr_vals[i] is None:
            continue
        if sma5[i] > sma10[i] and closes[i] > closes[i-1] and closes[i] > opens[i]:
            entry = closes[i]
            side = 'long'
            sl_price = entry - atr_vals[i] * params.sl_atr_multi
            tp_price = entry + atr_vals[i] * params.tp_atr_multi

            signals.append({
                'i': i, 'ts': klines['timestamps'][i], 'side': side,
                'entry': entry, 'sl': sl_price, 'tp': tp_price,
                'atr': atr_vals[i],
            })
            cooldown = params.cooldown_bars
    return signals


def signal_vwap_revert(klines: Dict, params) -> List[Dict]:
    """VWAP 回归"""
    closes = klines['closes']
    highs = klines['highs']
    lows = klines['lows']
    volumes = klines['volumes']
    atr_vals = atr(highs, lows, closes, 14)
    vwaps = vwap(closes, volumes)
    vwap_stds = vwap_std(closes, volumes, 20)

    signals = []
    cooldown = 0
    for i in range(30, len(closes)):
        cooldown -= 1
        if cooldown > 0:
            continue
        if vwaps[i] is None or vwap_stds[i] is None or atr_vals[i] is None:
            continue

        deviation = closes[i] - vwaps[i]
        std_dev = vwap_stds[i]

        # 量需 > 5日均量 × 1.2 (用前 120 bar 作 5日)
        if i < 121:
            continue
        avg_vol = sum(volumes[i-119:i+1]) / 120
        if volumes[i] < avg_vol * params.require_volume_multi:
            continue

        side = None
        if deviation < -std_dev * params.vwap_deviation_std:
            side = 'long'
        elif deviation > std_dev * params.vwap_deviation_std:
            side = 'short'
        if not side:
            continue

        entry = closes[i]
        # SL = entry ± 1σ (基于 VWAP std)
        sl_dist = std_dev * params.vwap_sl_std_multi
        sl_price = entry - sl_dist if side == 'long' else entry + sl_dist
        tp_price = vwaps[i] * (1 - params.vwap_tp_touch_pct/100) if side == 'long' else vwaps[i] * (1 + params.vwap_tp_touch_pct/100)

        signals.append({
            'i': i, 'ts': klines['timestamps'][i], 'side': side,
            'entry': entry, 'sl': sl_price, 'tp': tp_price,
            'atr': atr_vals[i],
        })
        cooldown = params.cooldown_bars
    return signals


def signal_turtle_breakout(klines: Dict, params) -> List[Dict]:
    """海龟通道突破"""
    closes = klines['closes']
    highs = klines['highs']
    lows = klines['lows']
    atr_vals = atr(highs, lows, closes, 14)

    hh, ll = donchian_breakout(highs, lows, params.turtle_channel_period)
    hh_exit, ll_exit = donchian_breakout(highs, lows, params.turtle_exit_channel_period)

    signals = []
    cooldown = 0
    for i in range(params.turtle_channel_period, len(closes)):
        cooldown -= 1
        if cooldown > 0:
            continue
        if hh[i-1] is None or ll[i-1] is None or atr_vals[i] is None:
            continue

        side = None
        if closes[i] > hh[i-1]:
            side = 'long'
        elif closes[i] < ll[i-1]:
            side = 'short'
        if not side:
            continue

        entry = closes[i]
        sl_price = entry - atr_vals[i] * params.sl_atr_multi if side == 'long' else entry + atr_vals[i] * params.sl_atr_multi
        tp_price = entry + atr_vals[i] * params.tp_atr_multi if side == 'long' else entry - atr_vals[i] * params.tp_atr_multi

        signals.append({
            'i': i, 'ts': klines['timestamps'][i], 'side': side,
            'entry': entry, 'sl': sl_price, 'tp': tp_price,
            'atr': atr_vals[i],
        })
        cooldown = params.cooldown_bars
    return signals


def signal_early_bird(klines: Dict, params) -> List[Dict]:
    """早盘动量: 假设 K 线是 5min, 开盘 30 min = 6 根 K 线
    看开盘 6 根 K 线的累计涨跌幅 + 量能
    """
    closes = klines['closes']
    opens = klines['opens']
    highs = klines['highs']
    lows = klines['lows']
    volumes = klines['volumes']
    atr_vals = atr(highs, lows, closes, 14)

    signals = []
    cooldown = 0
    # 简化: 找每根 K 线, 看 close vs 开盘 (5 bar 前) 的涨跌幅
    for i in range(20, len(closes) - params.max_hold_bars - 6):
        cooldown -= 1
        if cooldown > 0:
            continue

        # 取开盘 6 根 (5min × 6 = 30 min) 的累计涨跌
        open_price = opens[i - 5]  # 6 根前开 (第 1 根的开)
        window_high = max(highs[i-5:i+1])
        window_low = min(lows[i-5:i+1])
        window_vol = sum(volumes[i-5:i+1])

        # 跳空
        gap_pct = abs(opens[i] - closes[i-6]) / closes[i-6] * 100
        if gap_pct < params.early_bird_min_move_pct:
            continue

        # 量能
        if i < 121:
            continue
        avg_vol = sum(volumes[i-119:i+1]) / 120
        if window_vol < avg_vol * params.early_bird_volume_multi:
            continue

        # 顺势
        side = 'long' if closes[i] > opens[i] else 'short'
        entry = closes[i]
        sl_price = entry - atr_vals[i] * params.sl_atr_multi if side == 'long' else entry + atr_vals[i] * params.sl_atr_multi
        tp_price = entry + atr_vals[i] * params.tp_atr_multi if side == 'long' else entry - atr_vals[i] * params.tp_atr_multi

        signals.append({
            'i': i, 'ts': klines['timestamps'][i], 'side': side,
            'entry': entry, 'sl': sl_price, 'tp': tp_price,
            'atr': atr_vals[i],
        })
        cooldown = params.cooldown_bars
    return signals


SIGNAL_FNS = {
    'rsi2_revert': signal_rsi2_revert,
    'vwap_revert': signal_vwap_revert,
    'early_bird': signal_early_bird,
    'turtle_breakout': signal_turtle_breakout,
    'sma_breakout': signal_sma_breakout,
}


# ============== 回测执行 ==============

def run_backtest(klines: Dict, signals: List[Dict], symbol: str, strategy_name: str) -> Dict:
    """根据信号做回测.
    入场: 信号触发 (i 时刻 close)
    出场: SL / TP / max_hold_bars 三选一先到
    """
    closes = klines['closes']
    highs = klines['highs']
    lows = klines['lows']
    trades = []
    in_position = None  # {i_entry, side, entry, sl, tp}

    # 简化: 同时只能持 1 仓 (同向多仓不重入)
    for i in range(50, len(closes)):
        # 1) 平仓检查
        if in_position is not None:
            exit_price = None
            exit_reason = None
            i_entry = in_position['i_entry']
            side = in_position['side']
            sl = in_position['sl']
            tp = in_position['tp']

            # SL hit (用 high/low 检查)
            if side == 'long' and lows[i] <= sl:
                exit_price = sl
                exit_reason = 'SL'
            elif side == 'short' and highs[i] >= sl:
                exit_price = sl
                exit_reason = 'SL'
            elif side == 'long' and highs[i] >= tp:
                exit_price = tp
                exit_reason = 'TP'
            elif side == 'short' and lows[i] <= tp:
                exit_price = tp
                exit_reason = 'TP'
            elif i - i_entry >= 78:  # 默认 max_hold_bars
                exit_price = closes[i]
                exit_reason = 'EXPIRE'

            if exit_price is not None:
                pnl_pct = (exit_price - in_position['entry']) / in_position['entry'] * 100
                if side == 'short':
                    pnl_pct = -pnl_pct
                trades.append({
                    'side': side, 'entry': in_position['entry'], 'exit': exit_price,
                    'pnl_pct': pnl_pct, 'reason': exit_reason,
                    'i_entry': i_entry, 'i_exit': i,
                })
                in_position = None

        # 2) 入场检查
        for sig in signals:
            if sig['i'] == i and in_position is None:
                in_position = {
                    'i_entry': i, 'side': sig['side'],
                    'entry': sig['entry'], 'sl': sig['sl'], 'tp': sig['tp'],
                }
                break

    # 计算统计
    if not trades:
        return {
            'strategy': strategy_name, 'symbol': symbol,
            'signals': len(signals), 'trades': 0,
            'win_rate': 0, 'avg_pnl': 0, 'total_pnl': 0,
            'max_drawdown': 0, 'sharpe': 0,
        }

    wins = [t for t in trades if t['pnl_pct'] > 0]
    losses = [t for t in trades if t['pnl_pct'] <= 0]
    pnls = [t['pnl_pct'] for t in trades]
    win_rate = len(wins) / len(trades) * 100

    # 最大回撤 (累计收益曲线的 max drawdown)
    cum = [0]
    for p in pnls:
        cum.append(cum[-1] + p)
    peak = cum[0]
    max_dd = 0
    for v in cum:
        if v > peak:
            peak = v
        max_dd = min(max_dd, v - peak)

    # Sharpe 简化: 平均 / std
    avg = sum(pnls) / len(pnls)
    var = sum((x - avg)**2 for x in pnls) / len(pnls)
    std = var ** 0.5
    sharpe = avg / std if std > 0 else 0

    return {
        'strategy': strategy_name, 'symbol': symbol,
        'signals': len(signals), 'trades': len(trades),
        'wins': len(wins), 'losses': len(losses),
        'win_rate': round(win_rate, 1),
        'avg_pnl': round(avg, 3),
        'best': round(max(pnls), 2),
        'worst': round(min(pnls), 2),
        'total_pnl': round(sum(pnls), 2),
        'max_drawdown': round(max_dd, 2),
        'sharpe': round(sharpe, 2),
        'trades_detail': trades[:10],
    }


def fmt(result: Dict) -> str:
    """格式化回测报告"""
    lines = []
    lines.append(f"📊 {result['strategy']} {result['symbol']}")
    lines.append(f"  信号: {result['signals']} | 成交: {result['trades']} (W={result.get('wins',0)}, L={result.get('losses',0)})")
    if result['trades'] == 0:
        lines.append(f"  ⚠️ 无成交 (参数过严或市场平静)")
        return '\n'.join(lines)
    lines.append(f"  胜率: {result['win_rate']}%")
    lines.append(f"  平均盈亏: {result['avg_pnl']:+.3f}% | 最大盈: {result['best']:+.2f}% / 最大亏: {result['worst']:+.2f}%")
    lines.append(f"  累计盈亏: {result['total_pnl']:+.2f}% | 最大回撤: {result['max_drawdown']:+.2f}%")
    lines.append(f"  Sharpe: {result['sharpe']}")
    if result['trades'] > 0:
        lines.append(f"  最近 5 笔: {result['trades_detail'][:5]}")
    return '\n'.join(lines)


# ============== 主入口 ==============

def main():
    parser = argparse.ArgumentParser(description='策略回测 - v0.1')
    parser.add_argument('--strategy', choices=list(SIGNAL_FNS.keys()), required=True)
    parser.add_argument('--symbol', default='NVDA', help='Yahoo Finance ticker, e.g. NVDA / 0700.HK')
    parser.add_argument('--days', type=int, default=30)
    parser.add_argument('--interval', default='1h', help='K 线周期: 1h / 30m / 15m / 5m')
    args = parser.parse_args()

    print(f"⏳ 拉 {args.symbol} 最近 {args.days} 天 {args.interval} K线...")
    klines = fetch_klines(args.symbol, days=args.days, interval=args.interval)
    if not klines:
        print(f"❌ {args.symbol} 数据拉取失败")
        sys.exit(1)
    n = len(klines['closes'])
    print(f"✅ {n} 根 K 线")

    params = get_strategy(args.strategy)
    print(f"\n🎯 策略: {args.strategy}")
    print(f"📋 {params.name} (sl_atr={params.sl_atr_multi}, tp_atr={params.tp_atr_multi}, position={params.position_pct}%)")

    signals = SIGNAL_FNS[args.strategy](klines, params)
    print(f"🔍 信号数: {len(signals)}")

    result = run_backtest(klines, signals, args.symbol, args.strategy)
    print("\n" + fmt(result))

    # 输出 JSON
    result.pop('trades_detail', None)
    print(f"\n📊 JSON: {json.dumps(result, default=str, ensure_ascii=False)}")


if __name__ == '__main__':
    main()